from pathlib import Path
import shutil
import unittest
import random

import pandas as pd

from .utils import set_up_project

SAMPLE_FRAME = Path(__file__).parent / "sample_frame.jpg"
SAMPLE_FRAME_INPUT = Path(__file__).parent / "sample_frame_input.csv"
SAMPLE_AUTOCORRECT_OUTPUT = Path(__file__).parent / "sample_autocorrect_output.csv"


class Test2DTrialProcess(unittest.TestCase):
    """Test function performance on an actual trial - 2D, combined trial workflow"""

    def setUp(self):
        """Create trial with test data"""
        self.working_dir = Path.cwd() / "tmp"
        self.trial_csv, self.deepxromm = set_up_project(self.working_dir, "2D")
        self.deepxromm.xma_to_dlc()

    def test_first_frame_matches_in_dlc_csv(self):
        """When I run xma_to_dlc, does the DLC CSV have the same data as my original file?"""
        xmalab_data = pd.read_csv(self.trial_csv)
        xmalab_first_row = xmalab_data.loc[0, :]

        dlc_config = self.deepxromm.project.dlc_config.path_config_file
        labeled_data_path = dlc_config.parent / "labeled-data/MyData"
        dlc_data = pd.read_hdf(labeled_data_path / "CollectedData_NA.h5")
        cam1_img_path = str(
            (labeled_data_path / "test_cam1_0001.png").relative_to(dlc_config.parent)
        )
        cam2_img_path = str(
            (labeled_data_path / "test_cam2_0001.png").relative_to(dlc_config.parent)
        )
        cam1_first_row = dlc_data.loc[cam1_img_path, :]
        cam2_first_row = dlc_data.loc[cam2_img_path, :]
        for val in cam1_first_row.index:
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = xmalab_first_row[xmalab_key]
            dlc_data_point = cam1_first_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        for val in cam2_first_row.index:
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = xmalab_first_row[xmalab_key]
            dlc_data_point = cam2_first_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def test_last_frame_matches_in_dlc_csv(self):
        """When I run xma_to_dlc, does the DLC CSV have the same data as my original file?"""
        # Load XMAlab data
        xmalab_data = pd.read_csv(self.trial_csv)

        # Load DLC data
        dlc_config = self.deepxromm.project.dlc_config.path_config_file
        labeled_data_path = dlc_config.parent / "labeled-data/MyData"
        dlc_data = pd.read_hdf(labeled_data_path / "CollectedData_NA.h5")

        # Determine last frame included in training set
        last_file = Path(dlc_data.index[-1])
        last_frame_number = last_file.stem.split("_")[-1]
        last_frame_int = int(last_frame_number)

        # Load XMAlab last row
        xmalab_last_row = xmalab_data.loc[last_frame_int - 1]

        # Load DLC cam1 last row
        cam1_img_path = str(
            (labeled_data_path / f"test_cam1_{last_frame_number}.png").relative_to(
                dlc_config.parent
            )
        )
        cam1_last_row = dlc_data.loc[cam1_img_path, :]

        # Load DLC cam2 last row
        cam2_img_path = str(
            (labeled_data_path / f"test_cam2_{last_frame_number}.png").relative_to(
                dlc_config.parent
            )
        )
        cam2_last_row = dlc_data.loc[cam2_img_path, :]

        for val in cam1_last_row.index:
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = xmalab_last_row[xmalab_key]
            dlc_data_point = cam1_last_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        for val in cam2_last_row.index:
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = xmalab_last_row[xmalab_key]
            dlc_data_point = cam2_last_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def test_random_frame_matches_in_dlc_csv(self):
        """When I run xma_to_dlc, does the DLC CSV have the same data as my original file?"""
        # Load XMAlab data
        xmalab_data = pd.read_csv(self.trial_csv)

        # Load DLC data
        dlc_config = self.deepxromm.project.dlc_config.path_config_file
        labeled_data_path = dlc_config.parent / "labeled-data/MyData"
        dlc_data = pd.read_hdf(labeled_data_path / "CollectedData_NA.h5")

        # Determine rand frame included in training set
        file = Path(random.choice(dlc_data.index))
        frame_number = file.stem.split("_")[-1]
        frame_int = int(frame_number)

        # Load XMAlab rand row
        xmalab_row = xmalab_data.loc[frame_int - 1]

        # Load DLC cam1 rand row
        cam1_img_path = str(
            (labeled_data_path / f"test_cam1_{frame_number}.png").relative_to(
                dlc_config.parent
            )
        )
        cam1_row = dlc_data.loc[cam1_img_path, :]

        # Load DLC cam2 rand row
        cam2_img_path = str(
            (labeled_data_path / f"test_cam2_{frame_number}.png").relative_to(
                dlc_config.parent
            )
        )
        cam2_row = dlc_data.loc[cam2_img_path, :]

        for val in cam1_row.index:
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = xmalab_row[xmalab_key]
            dlc_data_point = cam1_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        for val in cam2_row.index:
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = xmalab_row[xmalab_key]
            dlc_data_point = cam2_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def test_create_training_dataset_succeeds(self):
        """Test that this project will create a dataset correctly"""
        self.deepxromm.create_training_dataset()

    def tearDown(self):
        """Remove the created temp project"""
        project_path = Path.cwd() / "tmp"
        shutil.rmtree(project_path)


class TestDlcToXma2D(unittest.TestCase):
    """Test dlc_to_xma function in 2D mode with round-trip verification"""

    def setUp(self):
        """Create 2D project and generate mock DLC analysis output"""
        self.working_dir = Path.cwd() / "tmp"
        _, self.deepxromm = set_up_project(self.working_dir, "2D")
        self.deepxromm.xma_to_dlc()
        self.mock_cam1_h5, self.mock_cam2_h5 = _copy_mock_dlc_data(self.working_dir)
        self.deepxromm.dlc_to_xma()

    def test_dlc_to_xma_creates_xmalab_format_files(self):
        """
        Given I have DLC analysis output with likelihood columns
        When I call dlc_to_xma
        Then it creates both CSV and HDF5 files
        """
        # Run dlc_to_xma
        output_dir = self.working_dir / "trials/test/it0"

        # Verify outputs exist
        xma_csv = output_dir / "test-Predicted2DPoints.csv"
        xma_h5 = output_dir / "test-Predicted2DPoints.h5"

        self.assertTrue(xma_csv.exists(), "XMAlab CSV not created")
        self.assertTrue(xma_h5.exists(), "XMAlab HDF5 not created")

    def test_dlc_to_xma_removes_likelihood_columns(self):
        """
        Given DLC analysis output has likelihood columns
        When I call dlc_to_xma
        Then the output XMAlab CSV has no likelihood columns
        And only contains X,Y coordinates
        """
        # Load and convert (split by camera)
        output_dir = self.working_dir / "trials/test/it0"

        # Check output
        xma_csv = output_dir / "test-Predicted2DPoints.csv"
        df = pd.read_csv(xma_csv)

        # Verify no likelihood columns
        for col in df.columns:
            self.assertNotIn("likelihood", col.lower())
            self.assertTrue(col.endswith(("_X", "_Y")))

    def test_first_frame_matches(self):
        """
        Given DLC-formatted output from a training run of this trial
        When I run dlc_to_xma
        Then the first frame of the reconstructed XMA data matches the DLC-formatted data
        """
        output_dir = self.working_dir / "trials/test/it0"
        cam1_dlc_data = pd.read_hdf(self.mock_cam1_h5)
        cam2_dlc_data = pd.read_hdf(self.mock_cam2_h5)

        xmalab_data = pd.read_csv(output_dir / "test-Predicted2DPoints.csv")
        xmalab_first_row = xmalab_data.loc[0, :]

        cam1_first_row = cam1_dlc_data.loc[0, :]
        for val in cam1_first_row.index:
            if "likelihood" in val or "marker001" not in val or "x" not in val:
                continue
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = xmalab_first_row[xmalab_key]
            dlc_data_point = cam1_first_row[val]
            with self.subTest(folder=val):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        cam2_first_row = cam2_dlc_data.loc[0, :]
        for val in cam2_first_row.index:
            if "likelihood" in val or "marker001" not in val or "x" not in val:
                continue
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = xmalab_first_row[xmalab_key]
            dlc_data_point = cam2_first_row[val]
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def test_last_frame_matches(self):
        """
        Given DLC-formatted output from a training run of this trial
        When I run dlc_to_xma
        Then the last frame of the reconstructed XMA data matches the DLC-formatted data
        """
        output_dir = self.working_dir / "trials/test/it0"
        cam1_dlc_data = pd.read_hdf(
            output_dir / "test_cam1DLC_resnet50_test_projectDec1shuffle1_100000.h5"
        )
        cam2_dlc_data = pd.read_hdf(
            output_dir / "test_cam2DLC_resnet50_test_projectDec1shuffle1_100000.h5"
        )

        xmalab_data = pd.read_csv(output_dir / "test-Predicted2DPoints.csv")
        # Find XMAlab last row
        xmalab_last_row_int = xmalab_data.index[-1]
        xmalab_last_row = xmalab_data.loc[xmalab_last_row_int, :]

        cam1_last_row = cam1_dlc_data.loc[xmalab_last_row_int, :]
        cam2_last_row = cam2_dlc_data.loc[xmalab_last_row_int, :]
        for val in cam1_last_row.index:
            if "likelihood" in val:
                continue
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = round(xmalab_last_row[xmalab_key], 4)
            dlc_data_point = round(cam1_last_row[val], 4)
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        for val in cam2_last_row.index:
            if "likelihood" in val:
                continue
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = round(xmalab_last_row[xmalab_key], 4)
            dlc_data_point = round(cam2_last_row[val], 4)
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def test_random_frame_matches(self):
        """
        Given DLC-formatted output from a training run of this trial
        When I run dlc_to_xma
        Then a random frame of the reconstructed XMA data matches the DLC-formatted data
        """
        output_dir = self.working_dir / "trials/test/it0"
        cam1_dlc_data = pd.read_hdf(
            output_dir / "test_cam1DLC_resnet50_test_projectDec1shuffle1_100000.h5"
        )
        cam2_dlc_data = pd.read_hdf(
            output_dir / "test_cam2DLC_resnet50_test_projectDec1shuffle1_100000.h5"
        )

        xmalab_data = pd.read_csv(output_dir / "test-Predicted2DPoints.csv")
        # Get random row from xmalab data
        xmalab_rand_row_int = random.choice(xmalab_data.index)
        xmalab_rand_row = xmalab_data.loc[xmalab_rand_row_int, :]

        cam1_rand_row = cam1_dlc_data.loc[xmalab_rand_row_int, :]
        cam2_rand_row = cam2_dlc_data.loc[xmalab_rand_row_int, :]
        for val in cam1_rand_row.index:
            if "likelihood" in val:
                continue
            xmalab_key = f"{val[1]}_cam1_{val[2].upper()}"
            xmalab_data_point = round(xmalab_rand_row[xmalab_key], 4)
            dlc_data_point = round(cam1_rand_row[val], 4)
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

        for val in cam2_rand_row.index:
            if "likelihood" in val:
                continue
            xmalab_key = f"{val[1]}_cam2_{val[2].upper()}"
            xmalab_data_point = round(xmalab_rand_row[xmalab_key], 4)
            dlc_data_point = round(cam2_rand_row[val], 4)
            with self.subTest(folder=xmalab_key):
                self.assertTrue(xmalab_data_point == dlc_data_point)

    def tearDown(self):
        """Remove the created temp project"""
        if self.working_dir.exists():
            shutil.rmtree(self.working_dir)


class TestAutocorrect2D(unittest.TestCase):
    """Test that autocorrect runs properly for 2D project"""

    def setUp(self):
        """Create 2D project and generate mock DLC analysis output"""
        self.working_dir = Path.cwd() / "tmp"
        self.working_dir = Path.cwd() / "tmp"
        _, self.deepxromm = set_up_project(self.working_dir, "2D")
        self.deepxromm.xma_to_dlc()
        self.mock_cam1_h5, self.mock_cam2_h5 = _copy_mock_dlc_data(self.working_dir)

    def run_autocorrect(self):
        """Run autocorrect using the provided deepxromm project"""
        self.deepxromm.autocorrect_trials()

    def tearDown(self):
        """Remove the created temp project"""
        if self.working_dir.exists():
            shutil.rmtree(self.working_dir)


def _copy_mock_dlc_data(project_dir: Path) -> tuple[Path, Path]:
    cam1_df = pd.read_hdf("trial_cam1dlc.h5")
    cam2_df = pd.read_hdf("trial_cam2dlc.h5")
    output_dir = project_dir / "trials/test/it0"
    output_dir.mkdir(parents=True, exist_ok=True)
    mock_cam1_h5 = (
        output_dir / "test_cam1DLC_resnet50_test_projectDec1shuffle1_100000.h5"
    )
    mock_cam2_h5 = (
        output_dir / "test_cam2DLC_resnet50_test_projectDec1shuffle1_100000.h5"
    )
    mock_cam1_csv = (
        output_dir / "test_cam1DLC_resnet50_test_projectDec1shuffle1_100000.csv"
    )
    mock_cam2_csv = (
        output_dir / "test_cam2DLC_resnet50_test_projectDec1shuffle1_100000.csv"
    )
    cam1_df.to_hdf(mock_cam1_h5, key="df_with_missing", mode="w")
    cam2_df.to_hdf(mock_cam2_h5, key="df_with_missing", mode="w")
    cam1_df.to_csv(mock_cam1_csv, na_rep="NaN")
    cam2_df.to_csv(mock_cam2_csv, na_rep="NaN")

    return mock_cam1_h5, mock_cam2_h5
