__all__ = ["analyze_latent"]

import re
from pathlib import Path
from typing import Sized, cast

import hydra
import numpy as np
import pandas as pd
import torch
from sklearn.base import TransformerMixin

import move.visualization as viz
from move.analysis.metrics import (
    calculate_accuracy,
    calculate_cosine_similarity,
)
from move.conf.schema import AnalyzeLatentConfig, MOVEConfig
from move.core.logging import get_logger
from move.core.typing import FloatArray
from move.data import io
from move.data.dataloaders import MOVEDataset, make_dataloader
from move.data.perturbations import (
    perturb_categorical_data,
    perturb_continuous_data,
)
from move.data.preprocessing import one_hot_encode_single
from move.models.vae import VAE
from move.training.training_loop import TrainingLoopOutput


def find_feature_values(
    feature_name: str,
    feature_names_lists: list[list[str]],
    feature_values: list[FloatArray],
) -> tuple[int, FloatArray]:
    """Look for the feature in the list of datasets and returns its values.

    Args:
        feature_name: Look-up key
        feature_names_lists: List of lists with feature names for each dataset
        feature_values: List of data arrays, each representing a dataset

    Raises:
        KeyError: If feature does not exist in any dataset

    Returns:
        Tuple containing (1) index of dataset containing feature and (2)
        values corresponding to the feature
    """
    _dataset_index, feature_index = [None] * 2
    for _dataset_index, feature_names in enumerate(feature_names_lists):
        try:
            feature_index = feature_names.index(feature_name)
        except ValueError:
            continue
        break
    if _dataset_index is not None and feature_index is not None:
        return (
            _dataset_index,
            np.take(feature_values[_dataset_index], feature_index, axis=1),
        )
    raise KeyError(f"Feature '{feature_name}' not in any dataset.")


def _validate_task_config(task_config: AnalyzeLatentConfig) -> None:
    if "_target_" not in task_config.reducer:
        raise ValueError("Reducer class not specified properly.")


def analyze_latent(config: MOVEConfig) -> None:
    """Train one model to inspect its latent space projections."""

    logger = get_logger(__name__)
    logger.info("Beginning task: analyze latent space")
    task_config = cast(AnalyzeLatentConfig, config.task)
    _validate_task_config(task_config)

    raw_data_path = Path(config.data.raw_data_path)
    interim_path = Path(config.data.interim_data_path)
    output_path = Path(config.data.results_path) / "latent_space"
    output_path.mkdir(exist_ok=True, parents=True)

    logger.debug("Reading data")
    sample_names = io.read_names(raw_data_path / f"{config.data.sample_names}.txt")
    cat_list, cat_names, con_list, con_names = io.load_preprocessed_data(
        interim_path,
        config.data.categorical_names,
        config.data.continuous_names,
    )
    test_dataloader = make_dataloader(
        cat_list,
        con_list,
        shuffle=False,
        batch_size=task_config.batch_size,
    )
    test_dataset = cast(MOVEDataset, test_dataloader.dataset)
    df_index = pd.Index(sample_names, name="sample")

    assert task_config.model is not None
    device = torch.device("cuda" if task_config.model.cuda else "cpu")
    model: VAE = hydra.utils.instantiate(
        task_config.model,
        continuous_shapes=test_dataset.con_shapes,
        categorical_shapes=test_dataset.cat_shapes,
    )

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

    model_path = output_path / "model.pt"
    if model_path.exists():
        logger.debug("Re-loading model")
        model.load_state_dict(torch.load(model_path))
        model.to(device)
    else:
        logger.debug("Training model")

        model.to(device)
        train_dataloader = make_dataloader(
            cat_list,
            con_list,
            shuffle=True,
            batch_size=task_config.batch_size,
            drop_last=True,
        )
        output: TrainingLoopOutput = hydra.utils.call(
            task_config.training_loop,
            model=model,
            train_dataloader=train_dataloader,
        )
        losses = output[:-1]
        torch.save(model.state_dict(), model_path)
        logger.info("Generating visualizations")
        logger.debug("Generating plot: loss curves")
        fig = viz.plot_loss_curves(losses)
        fig_path = str(output_path / "loss_curve.png")
        fig.savefig(fig_path, bbox_inches="tight")
        fig_df = pd.DataFrame(dict(zip(viz.LOSS_LABELS, losses)))
        fig_df.index.name = "epoch"
        fig_df.to_csv(output_path / "loss_curve.tsv", sep="\t")

    model.eval()

    logger.info("Projecting into latent space")
    latent_space = model.project(test_dataloader)
    reducer: TransformerMixin = hydra.utils.instantiate(task_config.reducer)
    embedding = reducer.fit_transform(latent_space)

    mappings_path = interim_path / "mappings.json"
    if mappings_path.exists():
        mappings = io.load_mappings(mappings_path)
    else:
        mappings = {}

    fig_df = pd.DataFrame(
        np.take(embedding, [0, 1], axis=1),
        columns=["dim0", "dim1"],
        index=df_index,
    )

    for feature_name in task_config.feature_names:
        logger.debug(f"Generating plot: latent space + '{feature_name}'")
        is_categorical = False
        try:
            dataset_index, feature_values = find_feature_values(
                feature_name, cat_names, cat_list
            )
            is_categorical = True
        except KeyError:
            try:
                dataset_index, feature_values = find_feature_values(
                    feature_name, con_names, con_list
                )
            except KeyError:
                logger.warning(f"Feature '{feature_name}' not found in any dataset.")
                continue

        if is_categorical:
            # Convert one-hot encoding to category codes
            is_nan = feature_values.sum(axis=1) == 0
            feature_values = np.argmax(feature_values, axis=1)

            dataset_name = config.data.categorical_names[dataset_index]
            feature_mapping = {
                str(code): category for category, code in mappings[dataset_name].items()
            }
            fig = viz.plot_latent_space_with_cat(
                embedding,
                feature_name,
                feature_values,
                feature_mapping,
                is_nan,
            )
            fig_df[feature_name] = np.where(is_nan, np.nan, feature_values)
        else:
            feature_values = feature_values
            fig = viz.plot_latent_space_with_con(
                embedding, feature_name, feature_values
            )
            fig_df[feature_name] = np.where(feature_values == 0, np.nan, feature_values)

        # Remove non-alpha characters
        safe_feature_name = re.sub(r"[^\w\s]", "", feature_name)
        fig_path = str(output_path / f"latent_space_{safe_feature_name}.png")
        fig.savefig(fig_path, bbox_inches="tight")

    fig_df.to_csv(output_path / "latent_space.tsv", sep="\t")

    logger.info("Reconstructing")
    cat_recons, con_recons = model.reconstruct(test_dataloader)
    con_recons = np.split(con_recons, np.cumsum(model.continuous_shapes[:-1]), axis=1)
    logger.info("Computing reconstruction metrics")
    scores = []
    labels = config.data.categorical_names + config.data.continuous_names
    for cat, cat_recon in zip(cat_list, cat_recons):
        accuracy = calculate_accuracy(cat, cat_recon)
        scores.append(accuracy)
    for con, con_recon in zip(con_list, con_recons):
        cosine_sim = calculate_cosine_similarity(con, con_recon)
        scores.append(cosine_sim)

    logger.debug("Generating plot: reconstruction metrics")

    plot_scores = [np.ma.compressed(np.ma.masked_equal(each, 0)) for each in scores]
    fig = viz.plot_metrics_boxplot(plot_scores, labels)
    fig_path = str(output_path / "reconstruction_metrics.png")
    fig.savefig(fig_path, bbox_inches="tight")
    fig_df = pd.DataFrame(dict(zip(labels, scores)), index=df_index)
    fig_df.to_csv(output_path / "reconstruction_metrics.tsv", sep="\t")

    logger.info("Computing feature importance")
    num_samples = len(cast(Sized, test_dataloader.sampler))
    for i, dataset_name in enumerate(config.data.categorical_names):
        logger.debug(f"Generating plot: feature importance '{dataset_name}'")
        na_value = one_hot_encode_single(mappings[dataset_name], None)
        dataloaders = perturb_categorical_data(
            test_dataloader, config.data.categorical_names, dataset_name, na_value
        )
        num_features = len(dataloaders)
        z = model.project(test_dataloader)
        diffs = np.empty((num_samples, num_features))
        for j, dataloader in enumerate(dataloaders):
            z_perturb = model.project(dataloader)
            diffs[:, j] = np.sum(z_perturb - z, axis=1)
        feature_mapping = {
            str(code): category for category, code in mappings[dataset_name].items()
        }
        fig = viz.plot_categorical_feature_importance(
            diffs, cat_list[i], cat_names[i], feature_mapping
        )
        fig_path = str(output_path / f"feat_importance_{dataset_name}.png")
        fig.savefig(fig_path, bbox_inches="tight")
        fig_df = pd.DataFrame(diffs, columns=cat_names[i], index=df_index)
        fig_df.to_csv(output_path / f"feat_importance_{dataset_name}.tsv", sep="\t")

    for i, dataset_name in enumerate(config.data.continuous_names):
        logger.debug(f"Generating plot: feature importance '{dataset_name}'")
        dataloaders = perturb_continuous_data(
            test_dataloader, config.data.continuous_names, dataset_name, 0.0
        )
        num_features = len(dataloaders)
        z = model.project(test_dataloader)
        diffs = np.empty((num_samples, num_features))
        for j, dataloader in enumerate(dataloaders):
            z_perturb = model.project(dataloader)
            diffs[:, j] = np.sum(z_perturb - z, axis=1)
        fig = viz.plot_continuous_feature_importance(diffs, con_list[i], con_names[i])
        fig_path = str(output_path / f"feat_importance_{dataset_name}.png")
        fig.savefig(fig_path, bbox_inches="tight")
        fig_df = pd.DataFrame(diffs, columns=con_names[i], index=df_index)
        fig_df.to_csv(output_path / f"feat_importance_{dataset_name}.tsv", sep="\t")