# License: BSD 3 clause
"""
Tests related to ablation experiments.

:author: Michael Heilman (mheilman@ets.org)
:author: Nitin Madnani (nmadnani@ets.org)
:author: Dan Blanchard (dblanchard@ets.org)
:author: Aoife Cahill (acahill@ets.org)
"""

import csv
import json
import unittest

from skll.experiments import run_configuration
from skll.utils.testing import (
    config_dir,
    create_jsonlines_feature_files,
    fill_in_config_paths,
    output_dir,
    remove_jsonlines_feature_files,
    test_dir,
    train_dir,
    unlink,
)


class TestAblation(unittest.TestCase):
    """Test class for ablation tests."""

    @classmethod
    def setUpClass(cls):
        """Create necessary directories for testing."""
        for dir_path in [train_dir, test_dir, output_dir]:
            dir_path.mkdir(exist_ok=True)

        # create jsonlines feature files
        create_jsonlines_feature_files(train_dir)

    @classmethod
    def tearDownClass(cls):
        """Clean up after tests."""
        for output_file in output_dir.glob("ablation_cv_*"):
            unlink(output_file)

        config_files = [
            "test_ablation.cfg",
            "test_ablation_all_combos.cfg",
            "test_ablation_feature_hasher.cfg",
            "test_ablation_feature_hasher_all_combos.cfg",
            "test_ablation_sampler.cfg",
            "test_ablation_sampler_all_combos.cfg",
            "test_ablation_feature_hasher_sampler.cfg",
            "test_ablation_feature_hasher_sampler_all_combos.cfg",
        ]
        for cf in config_files:
            unlink(config_dir / cf)

        remove_jsonlines_feature_files(train_dir)

    def check_ablation_rows(self, reader):
        """
        Ensure that all ablated faetures and featureset values are correct.

        Parameters
        ----------
        reader : object
            Reader object that iterates over the output summary file.

        Returns
        -------
        int
            The total number of rows in the summary file.
        """
        row_num = 0
        for row_num, row in enumerate(reader, 1):
            if row["ablated_features"]:
                fs_str, ablated_str = row["featureset_name"].split("_minus_")
                actual_ablated = json.loads(row["ablated_features"])
            else:
                fs_str, ablated_str = row["featureset_name"].split("_all")
                actual_ablated = []
            expected_fs = set(fs_str.split("+"))
            expected_ablated = ablated_str.split("+") if ablated_str else []
            expected_fs = sorted(expected_fs - set(expected_ablated))
            actual_fs = json.loads(row["featureset"])
            self.assertEqual(expected_ablated, actual_ablated)
            self.assertEqual(expected_fs, actual_fs)
        return row_num

    def test_ablation_cv(self):
        """Test ablation + cross-validation."""
        config_template_path = config_dir / "test_ablation.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=1, local=True)

        # read in the summary file and make sure it has
        # 7 ablated featuresets * (10 folds + 1 average line) * 2 learners = 154
        # lines
        with open(output_dir / "ablation_cv_plain_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 154)

        # make sure there are 7 ablated featuresets * 2 learners = 12 results files
        num_result_files = len(list(output_dir.glob("ablation_cv_plain*.results")))
        self.assertEqual(num_result_files, 14)

    def test_ablation_cv_all_combos(self):
        """Test ablation all-combos + cross-validation."""
        config_template_path = config_dir / "test_ablation_all_combos.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=None, local=True)

        # read in the summary file and make sure it has
        # 10 ablated featuresets * (10 folds + 1 average line) * 2 learners = 220
        # lines
        with open(output_dir / "ablation_cv_plain_all_combos_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 220)

        # make sure there are 10 ablated featuresets * 2 learners = 20 results
        # files
        num_result_files = len(list(output_dir.glob("ablation_cv_plain_all_combos*results")))
        self.assertEqual(num_result_files, 20)

    def test_ablation_cv_feature_hasher(self):
        """Test ablation + cross-validation + feature hashing."""
        config_template_path = config_dir / "test_ablation_feature_hasher.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=1, local=True)

        # read in the summary file and make sure it has
        # 7 ablated featuresets * (10 folds + 1 average line) * 2 learners = 154
        # lines
        with open(output_dir / "ablation_cv_feature_hasher_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 154)

        # make sure there are 7 ablated featuresets * 2 learners = 14 results files
        num_result_files = len(list(output_dir.glob("ablation_cv_feature_hasher_*.results")))
        self.assertEqual(num_result_files, 14)

    def test_ablation_cv_feature_hasher_all_combos(self):
        """Test ablation all-combos + cross-validation + feature hashing."""
        config_template_path = config_dir / "test_ablation_feature_hasher_all_combos.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=None, local=True)

        # read in the summary file and make sure it has
        #    10 ablated featuresets
        #      * (10 folds + 1 average line)
        #      * 2 learners
        #    = 220 lines in total
        with open(output_dir / "ablation_cv_feature_hasher_all_combos_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 220)

        # make sure there are 10 ablated featuresets * 2 learners = 20 results
        # files
        num_result_files = len(
            list(output_dir.glob("ablation_cv_feature_hasher_all_combos*.results"))
        )
        self.assertEqual(num_result_files, 20)

    def test_ablation_cv_sampler(self):
        """Test ablation + cross-validation + samplers."""
        config_template_path = config_dir / "test_ablation_sampler.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=1, local=True)

        # read in the summary file and make sure it has
        # 7 ablated featuresets * (10 folds + 1 average line) * 2 learners = 154
        # lines
        with open(output_dir / "ablation_cv_sampler_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 154)

        # make sure there are 6 ablated featuresets * 2 learners = 12 results files
        num_result_files = len(list(output_dir.glob("ablation_cv_sampler_f*.results")))
        self.assertEqual(num_result_files, 14)

    def test_ablation_cv_all_combos_sampler(self):
        """Test ablation all-combos + cross-validation + samplers."""
        config_template_path = config_dir / "test_ablation_sampler_all_combos.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=None, local=True)

        # read in the summary file and make sure it has
        # 10 ablated featuresets * (10 folds + 1 average line) * 2 learners = 220
        # lines
        with open(output_dir / "ablation_cv_sampler_all_combos_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 220)

        # make sure there are 10 ablated featuresets * 2 learners = 20 results
        # files
        num_result_files = len(list(output_dir.glob("ablation_cv_sampler_all_combos*.results")))
        self.assertEqual(num_result_files, 20)

    def test_ablation_cv_feature_hasher_sampler(self):
        """Test ablation + cross-validation + feature hashing + samplers."""
        config_template_path = config_dir / "test_ablation_feature_hasher_sampler.template.cfg"
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=1, local=True)

        # read in the summary file and make sure it has
        # 7 ablated featuresets * (10 folds + 1 average line) * 2 learners = 154
        # lines
        with open(output_dir / "ablation_cv_feature_hasher_sampler_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 154)

        # make sure there are 7 ablated featuresets * 2 learners = 14 results files
        num_result_files = len(
            list(output_dir.glob("ablation_cv_feature_hasher_sampler_f*.results"))
        )
        self.assertEqual(num_result_files, 14)

    def test_ablation_cv_feature_hasher_all_combos_sampler(self):
        """Test ablation all-combos + cross-validation + feature hashing + samplers."""
        config_template_path = (
            config_dir / "test_ablation_feature_hasher_sampler_all_combos.template.cfg"
        )
        config_path = fill_in_config_paths(config_template_path)

        run_configuration(config_path, quiet=True, ablation=None, local=True)

        # read in the summary file and make sure it has
        # 10 ablated featuresets * (10 folds + 1 average line) * 2 learners = 220
        # lines
        with open(output_dir / "ablation_cv_feature_hasher_sampler_all_combos_summary.tsv") as f:
            reader = csv.DictReader(f, dialect=csv.excel_tab)
            num_rows = self.check_ablation_rows(reader)
            self.assertEqual(num_rows, 220)

        # make sure there are 10 ablated featuresets * 2 learners = 20 results
        # files
        num_result_files = len(
            list(output_dir.glob("ablation_cv_feature_hasher_sampler_all_combos*.results"))
        )
        self.assertEqual(num_result_files, 20)
