from dataflow.operators.code import (
    CodeAutoGeneratedFilter,
    CodeLengthSampleFilter,
    CodeTextCompositionFilter,
    CodeEncodedDataFilter,
    CodeDocumentQualityFilter,
    CodeFileTypeContentFilter,
    CodeGenericScoreFilter,
)
from dataflow.utils.storage import FileStorage

class PTCodeFilter_CPUPipeline():
    def __init__(self):
        self.storage = FileStorage(
            first_entry_file_name="../example_data/CodePipeline/code_input.jsonl",
            cache_path="./cache",
            file_name_prefix="dataflow_cache_step",
            cache_type="jsonl",
        )
        
        self.autogen_filter_step1 = CodeAutoGeneratedFilter(
            min_score=1.0,
            max_score=1.0
        )
        
        self.length_filter_step2 = CodeLengthSampleFilter(
            min_score=1.0,
            max_score=1.0
        )
        
        self.text_composition_filter_step3 = CodeTextCompositionFilter(
            min_score=1.0,
            max_score=1.0
        )
        
        self.encoded_data_filter_step4 = CodeEncodedDataFilter(
            min_score=1.0,
            max_score=1.0
        )
        
        self.doc_quality_filter_step5 = CodeDocumentQualityFilter(
            min_score=1.0,
            max_score=1.0,
            thresholds={
                'min_num_chars': 100,
                'max_num_chars': 100000,
                'min_num_words': 10,
                'max_num_words': 50000,
                'max_frac_duplicate_lines': 0.3,
                'max_frac_duplicate_2gram': 0.3,
                'max_frac_duplicate_3gram': 0.3,
                'max_frac_duplicate_4gram': 0.3,
                'max_frac_duplicate_5gram': 0.3,
                'max_frac_curly_bracket': 0.1,
                'max_frac_all_caps_words': 0.3,
                'min_entropy_unigram': 2.0,
            }
        )
        
        self.file_type_filter_step6 = CodeFileTypeContentFilter()
        
        self.score_filter_step7 = CodeGenericScoreFilter()
    
    def forward(self):
        self.autogen_filter_step1.run(
            storage=self.storage.step(),
            input_key="lines",
            output_key="autogen_filter_label"
        )
        
        self.length_filter_step2.run(
            storage=self.storage.step(),
            input_key="lines",
            output_key="length_filter_label"
        )
        
        self.text_composition_filter_step3.run(
            storage=self.storage.step(),
            input_key="text",
            output_key="text_composition_filter_label"
        )
        
        self.encoded_data_filter_step4.run(
            storage=self.storage.step(),
            input_key="text",
            output_key="encoded_data_filter_label"
        )
        
        self.doc_quality_filter_step5.run(
            storage=self.storage.step(),
            input_key="text",
            output_key="doc_quality_filter_label"
        )
        
        self.file_type_filter_step6.run(
            storage=self.storage.step(),
            input_key="dataframe",
            output_key="file_type_filter_label"
        )
        
        # self.score_filter_step7.run(
        #     storage=self.storage.step(),
        #     input_key="quality_score",
        #     output_key="score_filter_label",
        #     score_threshold=8,
        #     filter_method="greater_equal"
        # )

if __name__ == "__main__":
    model = PTCodeFilter_CPUPipeline()
    model.forward()