import numpy as np
import pandas as pd
from schools3.ml.experiments.models_experiment import ModelsExperiment
from schools3.data.base.cohort import Cohort
from schools3.config import main_config
from schools3.config import global_config
from schools3.data.datasets.dataset import Dataset
from schools3.ml.experiments.feat_importances_experiment import FeatureImportancesExperiment
from schools3.ml.experiments.single_dataset_experiment import SingleDatasetExperiment
from schools3.ml.models.tfkeras_model import TFKerasModel
from schools3.ml.models.sklearn_model import SklearnModel
import schools3.config.ml.experiments.feat_pruning_experiment_config as config
from schools3.config.data.datasets import dataset_config

# an experiment that trains models with subsets of the features according to their permutation importance rank
# like SingleDatasetExperiment, this works on a specific grade
class FeaturePruningExperiment(ModelsExperiment):
    def __init__(
        self, name='ignore',
        features_list=main_config.features,
        labels=main_config.labels,
        models=main_config.models,
        metrics=main_config.metrics,
        use_cache=main_config.use_cache
    ):
        super(FeaturePruningExperiment, self).__init__(
            name, features_list, labels, models, metrics, use_cache=use_cache
        )

    def perform(
        self, grade=main_config.single_grade,
        train_years=main_config.train_years,
        test_years=main_config.test_years,
        compute_train_metrics=False, **kwargs
    ):
        train_cohort = Cohort(grade, train_years)

        df = pd.DataFrame()
        for model in self.models:
            if not (isinstance(model, SklearnModel) or isinstance(model, TFKerasModel)):
                continue

            train_data = Dataset(train_cohort, self.features_list, model.get_feature_processor(), self.labels)
            model.train(train_data)

            feats_exp = FeatureImportancesExperiment('ignore', self.features_list, self.labels, [model], self.metrics)
            
            feature_names, _, sorted_idxs = feats_exp.get_feature_importances(model, train_data)
            feats = np.flip(feature_names[sorted_idxs])

            for i in config.num_feats:
                dataset_config.feat_whitelist.clear()
                for feat in feats[:i]:
                    dataset_config.feat_whitelist.append(feat)

                exp = SingleDatasetExperiment('ignore', self.features_list, self.labels, [model], self.metrics)
                cur_df = exp.perform(grade, train_years, test_years, compute_train_metrics=compute_train_metrics, **kwargs)
                cur_df['num_feats'] = i

                df = pd.concat([df, cur_df], ignore_index=True)

        return df

[((36, 23, 36, 49), 'schools3.data.base.cohort.Cohort', 'Cohort', ({(36, 30, 36, 35): 'grade', (36, 37, 36, 48): 'train_years'}, {}), '(grade, train_years)', False, 'from schools3.data.base.cohort import Cohort\n'), ((38, 13, 38, 27), 'pandas.DataFrame', 'pd.DataFrame', ({}, {}), '()', True, 'import pandas as pd\n'), ((46, 24, 46, 118), 'schools3.ml.experiments.feat_importances_experiment.FeatureImportancesExperiment', 'FeatureImportancesExperiment', ({(46, 53, 46, 61): '"""ignore"""', (46, 63, 46, 81): 'self.features_list', (46, 83, 46, 94): 'self.labels', (46, 96, 46, 103): '[model]', (46, 105, 46, 117): 'self.metrics'}, {}), "('ignore', self.features_list, self.labels, [\n    model], self.metrics)", False, 'from schools3.ml.experiments.feat_importances_experiment import FeatureImportancesExperiment\n'), ((49, 20, 49, 55), 'numpy.flip', 'np.flip', ({(49, 28, 49, 54): 'feature_names[sorted_idxs]'}, {}), '(feature_names[sorted_idxs])', True, 'import numpy as np\n'), ((52, 16, 52, 53), 'schools3.config.data.datasets.dataset_config.feat_whitelist.clear', 'dataset_config.feat_whitelist.clear', ({}, {}), '()', False, 'from schools3.config.data.datasets import dataset_config\n'), ((56, 22, 56, 111), 'schools3.ml.experiments.single_dataset_experiment.SingleDatasetExperiment', 'SingleDatasetExperiment', ({(56, 46, 56, 54): '"""ignore"""', (56, 56, 56, 74): 'self.features_list', (56, 76, 56, 87): 'self.labels', (56, 89, 56, 96): '[model]', (56, 98, 56, 110): 'self.metrics'}, {}), "('ignore', self.features_list, self.labels, [model],\n    self.metrics)", False, 'from schools3.ml.experiments.single_dataset_experiment import SingleDatasetExperiment\n'), ((60, 21, 60, 63), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n'), ((54, 20, 54, 62), 'schools3.config.data.datasets.dataset_config.feat_whitelist.append', 'dataset_config.feat_whitelist.append', ({(54, 57, 54, 61): 'feat'}, {}), '(feat)', False, 'from schools3.config.data.datasets import dataset_config\n')]