__all__ = ["tune_model"]

from pathlib import Path

from typing import Any, Literal, cast

import hydra

import numpy as np

import pandas as pd

import torch

from hydra.core.hydra_config import HydraConfig

from hydra.types import RunMode

from matplotlib.cbook import boxplot_stats

from numpy.typing import ArrayLike

from omegaconf import OmegaConf

from sklearn.metrics.pairwise import cosine_similarity

from move.analysis.metrics import (

    calculate_accuracy,

    calculate_cosine_similarity,

)

from move.conf.schema import (

    MOVEConfig,

    TuneModelConfig,

    TuneModelReconstructionConfig,

    TuneModelStabilityConfig,

)

from move.core.logging import get_logger

from move.core.typing import BoolArray

from move.data import io

from move.data.dataloaders import MOVEDataset, make_dataloader, split_samples

from move.models.vae import VAE

TaskType = Literal["reconstruction", "stability"]

def _get_task_type(

    task_config: TuneModelConfig,

) -> TaskType:

    task_type = OmegaConf.get_type(task_config)

    if task_type is TuneModelReconstructionConfig:

        return "reconstruction"

    if task_type is TuneModelStabilityConfig:

        return "stability"

    raise ValueError("Unsupported type of task!")

def _get_record(values: ArrayLike, **kwargs) -> dict[str, Any]:

    record = kwargs

    bxp_stats, *_ = boxplot_stats(values)

    bxp_stats.pop("fliers")

    record.update(bxp_stats)

    return record

def tune_model(config: MOVEConfig) -> float:

    """Train multiple models to tune the model hyperparameters."""

    hydra_config = HydraConfig.get()

    if hydra_config.mode != RunMode.MULTIRUN:

        raise ValueError("This task must run in multirun mode.")

    # Delete sweep run config

    sweep_config_path = Path(hydra_config.sweep.dir).joinpath("multirun.yaml")

    if sweep_config_path.exists():

        sweep_config_path.unlink()

    job_num = hydra_config.job.num + 1

    logger = get_logger(__name__)

    task_config = cast(TuneModelConfig, config.task)

    task_type = _get_task_type(task_config)

    logger.info(f"Beginning task: tune model {task_type} {job_num}")

    logger.info(f"Job name: {hydra_config.job.override_dirname}")

    interim_path = Path(config.data.interim_data_path)

    output_path = Path(config.data.results_path) / "tune_model"

    output_path.mkdir(exist_ok=True, parents=True)

    logger.debug("Reading data")

    cat_list, _, con_list, _ = io.load_preprocessed_data(

        interim_path,

        config.data.categorical_names,

        config.data.continuous_names,

    )

    assert task_config.model is not None

    device = torch.device("cuda" if task_config.model.cuda is True else "cpu")

    def _tune_stability(

        task_config: TuneModelStabilityConfig,

    ):

        label = [hp.split("=") for hp in hydra_config.job.override_dirname.split(",")]

        train_dataloader = make_dataloader(

            cat_list,

            con_list,

            shuffle=True,

            batch_size=task_config.batch_size,

            drop_last=True,

        )

        test_dataloader = make_dataloader(

            cat_list,

            con_list,

            shuffle=False,

            batch_size=task_config.batch_size,

            drop_last=False,

        )

        train_dataset = cast(MOVEDataset, train_dataloader.dataset)

        logger.info(f"Training {task_config.num_refits} refits")

        cosine_sim0 = None

        cosine_sim_diffs = []

        for j in range(task_config.num_refits):

            logger.debug(f"Refit: {j + 1}/{task_config.num_refits}")

            model: VAE = hydra.utils.instantiate(

                task_config.model,

                continuous_shapes=train_dataset.con_shapes,

                categorical_shapes=train_dataset.cat_shapes,

            )

            model.to(device)

            hydra.utils.call(

                task_config.training_loop,

                model=model,

                train_dataloader=train_dataloader,

            )

            model.eval()

            latent, *_ = model.latent(test_dataloader, kld_weight=1)

            if cosine_sim0 is None:

                cosine_sim0 = cosine_similarity(latent)

            else:

                cosine_sim = cosine_similarity(latent)

                D = np.absolute(cosine_sim - cosine_sim0)

                # removing the diagonal element (cos_sim with itself)

                diff = D[~np.eye(D.shape[0], dtype=bool)].reshape(D.shape[0], -1)

                mean_diff = np.mean(diff)

                cosine_sim_diffs.append(mean_diff)

        record = _get_record(

            cosine_sim_diffs,

            job_num=job_num,

            **dict(label),

            metric="mean_diff_cosine_similarity",

            num_refits=task_config.num_refits,

        )

        logger.info("Writing results")

        df_path = output_path / "stability_stats.tsv"

        header = not df_path.exists()

        df = pd.DataFrame.from_records([record])

        df.to_csv(df_path, sep="\t", mode="a", header=header, index=False)

    def _tune_reconstruction(

        task_config: TuneModelReconstructionConfig,

    ):

        split_path = interim_path / "split_mask.npy"

        if split_path.exists():

            split_mask: BoolArray = np.load(split_path)

        else:

            num_samples = cat_list[0].shape[0] if cat_list else con_list[0].shape[0]

            split_mask = split_samples(num_samples, 0.9)

            np.save(split_path, split_mask)

        train_dataloader = make_dataloader(

            cat_list,

            con_list,

            split_mask,

            shuffle=True,

            batch_size=task_config.batch_size,

            drop_last=True,

        )

        train_dataset = cast(MOVEDataset, train_dataloader.dataset)

        model: VAE = hydra.utils.instantiate(

            task_config.model,

            continuous_shapes=train_dataset.con_shapes,

            categorical_shapes=train_dataset.cat_shapes,

        )

        model.to(device)

        logger.debug(f"Model: {model}")

        logger.debug("Training model")

        hydra.utils.call(

            task_config.training_loop,

            model=model,

            train_dataloader=train_dataloader,

        )

        model.eval()

        logger.info("Reconstructing")

        logger.info("Computing reconstruction metrics")

        label = [hp.split("=") for hp in hydra_config.job.override_dirname.split(";")]

        records = []

        splits = zip(["train", "test"], [split_mask, ~split_mask])

        for split_name, mask in splits:

            dataloader = make_dataloader(

                cat_list,

                con_list,

                mask,

                shuffle=False,

                batch_size=task_config.batch_size,

            )

            cat_recons, con_recons = model.reconstruct(dataloader)

            con_recons = np.split(

                con_recons, np.cumsum(model.continuous_shapes[:-1]), axis=1

            )

            for cat, cat_recon, dataset_name in zip(

                cat_list, cat_recons, config.data.categorical_names

            ):

                logger.debug(f"Computing accuracy: '{dataset_name}'")

                accuracy = calculate_accuracy(cat[mask], cat_recon)

                record = _get_record(

                    accuracy,

                    job_num=job_num,

                    **dict(label),

                    metric="accuracy",

                    dataset=dataset_name,

                    split=split_name,

                )

                records.append(record)

            for con, con_recon, dataset_name in zip(

                con_list, con_recons, config.data.continuous_names

            ):

                logger.debug(f"Computing cosine similarity: '{dataset_name}'")

                cosine_sim = calculate_cosine_similarity(con[mask], con_recon)

                record = _get_record(

                    cosine_sim,

                    job_num=job_num,

                    **dict(label),

                    metric="cosine_similarity",

                    dataset=dataset_name,

                    split=split_name,

                )

                records.append(record)

        logger.info("Writing results")

        df_path = output_path / "reconstruction_stats.tsv"

        header = not df_path.exists()

        df = pd.DataFrame.from_records(records)

        df.to_csv(df_path, sep="\t", mode="a", header=header, index=False)

    if task_type == "reconstruction":

        task_config = cast(TuneModelReconstructionConfig, task_config)

        _tune_reconstruction(task_config)

    elif task_type == "stability":

        task_config = cast(TuneModelStabilityConfig, task_config)

        _tune_stability(task_config)

    return 0.0