import ast
import astor
from nbminer.features.featurize.cell_features import SegmentFeatures


class ResampleByNode:
    def __init__(self):
        pass

    def transform(self, corpus):
        self.corpus = corpus
        for index in range(corpus.get_number_notebooks()):
            self.get_new_notebook(index)
        return self.corpus

    def get_new_notebook(self, index):
        new_order = []
        cell_count = 0
        for cell in self.corpus.get_notebook_segments(index):
            cell_count += 1
            if not cell.feature_exists('ast'):
                cell.put_feature('ast', ast.parse(cell.get_feature('code')))
            cell_ast = cell.get_feature('ast')
            for item in cell_ast.body:
                new_ast = ast.parse('')
                new_ast.body.append(item)
                segment = SegmentFeatures()
                segment.put_feature('ast', new_ast)
                segment.put_feature('original_code', cell.get_feature('code'))
                segment.put_feature('notebook_name', cell.get_feature('notebook_name'))
                segment.put_feature('cell_id', str(cell_count))
                try:
                    segment.put_feature('astor', astor.to_source(new_ast))
                except:
                    print(ast.dump(new_ast))
                new_order.append(segment)

        self.corpus.set_notebook_segments(index, new_order)
