import itertools
import json
from pathlib import Path

import pandas as pd

from mist.app.loggers.logger import logger


class MistDists:
    """
    Class for generating distance and allele matrices from typing outputs.
    Expected input:
    - a list of TSV or JSON files generated by MiST
    """

    MIN_NB_DATASETS = 3

    def __init__(
            self, inputs: list[Path], out_matrix: Path, out_dists: Path, min_perc_loci: int,
            min_perc_samples: int) -> None:
        """
        Initializes the main script.
        :param inputs: List of input files
        :param out_matrix: Output allele matrix
        :param out_dists: Output file with pairwise distances
        :param min_perc_loci: Remove datasets with <% of loci detected
        :param min_perc_samples: Remove loci detected in <% of samples
        :return: None
        """
        self._inputs = inputs
        self._out_matrix = out_matrix
        self._out_dists = out_dists
        self._min_perc_loci = min_perc_loci
        self._min_perc_samples = min_perc_samples

    @staticmethod
    def _calc_distance(alleles_a: pd.Series, alleles_b: pd.Series) -> int:
        """
        Calculate the pairwise allele distance.
        :param alleles_a: Alleles a
        :param alleles_b: Alleles b
        :return: Number of allelic differences
        """
        total_dist = 0
        for allele_a, allele_b in zip(alleles_a, alleles_b):
            if allele_a == '-' and allele_b == '-':
                continue
            else:
                total_dist += 0 if allele_a == allele_b else 1
        return total_dist

    @staticmethod
    def parse_tsv(path_in: Path) -> tuple[str, pd.Series]:
        """
        Parses the input TSV file.
        :param path_in: Input path
        :return: Parsed data
        """
        logger.debug(f'Parsing file: {path_in}')
        data_tsv = pd.read_table(path_in, dtype=str, index_col='locus')
        return path_in.name.replace('.tsv', ''), data_tsv['allele']

    @staticmethod
    def parse_json(path_in: Path) -> tuple[str, pd.Series]:
        """
        Parses the input JSON file.
        :param path_in: Input path
        :return: Parsed data
        """
        logger.debug(f'Parsing file: {path_in}')
        with path_in.open() as handle:
            data = json.load(handle)
        alleles = pd.Series({locus: row['allele_str'] for locus, row in data['alleles'].items()})
        return path_in.name.replace('.json', ''), alleles

    def _log_nb_perfect_hits(self, name, alleles: pd.Series) -> None:
        """
        Logs the number and percentage of perfect hits for the input dataset.
        :param name: Dataset name
        :param alleles: Allele calls
        :return: None
        """
        nb_perfect = len(alleles) - alleles.eq('-').sum()
        perc_perfect = 100 * nb_perfect / len(alleles)
        logger.info(f"{name}: {nb_perfect}/{len(alleles):,} ({perc_perfect:.2f}%) perfect hits")

    def _parse_input_files(self) -> pd.DataFrame:
        """
        Parse the input files.
        :return: DataFrame with detected alleles (index = sample name)
        """
        # Parse the input
        datasets = {}
        for path_in in self._inputs:
            if path_in.suffix == '.tsv':
                name, alleles = MistDists.parse_tsv(path_in)
                datasets[name] = alleles
            elif path_in.suffix == '.json':
                name, alleles = MistDists.parse_json(path_in)
                datasets[name] = alleles
            else:
                raise ValueError(f'Invalid extension: {path_in.suffix} (expected .tsv or .json)')
            self._log_nb_perfect_hits(name, alleles)
        if len(datasets) < MistDists.MIN_NB_DATASETS:
            raise ValueError(f'At least {MistDists.MIN_NB_DATASETS} input files are required (found: {len(datasets)})')

        # Create merged DataFrame
        return pd.DataFrame(data=list(datasets.values()), index=list(datasets.keys()), dtype=str)

    def _calc_distance_matrix(self, allele_data_filt: pd.DataFrame) -> pd.DataFrame:
        """
        Calculates the pairwise distance matrix
        :param allele_data_filt: Filtered allele matrix
        :return: Distance matrix
        """
        # Calculate pair-wise distances
        distance_by_dataset_pair = {}
        for dataset_a, dataset_b in itertools.combinations(allele_data_filt.index, r=2):
            key = tuple(sorted([dataset_a, dataset_b]))
            dist = MistDists._calc_distance(allele_data_filt.loc[dataset_a], allele_data_filt.loc[dataset_b])
            distance_by_dataset_pair[key] = dist

        # Create data frame with pairwise distances
        records_out = []
        for dataset_a in allele_data_filt.index:
            records_out.append({
                dataset_b: distance_by_dataset_pair.get(tuple(sorted([dataset_a, dataset_b])), 0) for
                dataset_b in allele_data_filt.index
            })
        return pd.DataFrame(records_out, index=allele_data_filt.index)

    def _filt_allele_matrix(self, allele_data: pd.DataFrame) -> pd.DataFrame:
        """
        Filters the allele matrix by removing:
        - Datasets with less than x% of loci detected
        - Loci present in less than x% of datasets
        :param allele_data: Allele data
        :return: Filtered allele data, loci cutoff, datasets cutoff
        """
        # Filter allele matrix (nb. of loci detected per dataset)
        nb_loci_detected = allele_data.apply(lambda x: len(x) - list(x).count('-'), axis=1)
        cutoff_loci = int(self._min_perc_loci * len(allele_data.columns) / 100)
        logger.info(f"Removing datasets with <{cutoff_loci} ({self._min_perc_loci}%) loci detected")
        allele_data_filt = allele_data[nb_loci_detected > cutoff_loci]
        logger.info(f"{len(allele_data_filt)}/{len(allele_data)} datasets passed filtering")

        # Filter allele matrix (loci detected in nb. of datasets)
        cutoff_datasets = int(self._min_perc_samples * len(allele_data_filt) / 100)
        logger.info(f"Removing loci detected in <{cutoff_datasets} ({self._min_perc_samples}%) samples")
        # noinspection PyUnresolvedReferences
        locus_present_in_datasets = (allele_data_filt != '-').sum(axis=0)
        allele_data_filt = allele_data_filt.loc[:, locus_present_in_datasets > cutoff_datasets]
        logger.info(f"{len(allele_data_filt.columns)}/{len(allele_data.columns)} loci passed filtering")
        return allele_data_filt

    def run(self) -> None:
        """
        The main method to construct the MLST phylogeny.
        :return: None
        """
        # Parse the input files
        df_alleles = self._parse_input_files()

        # Filter the allele matrix
        df_alleles_filt = self._filt_allele_matrix(df_alleles)
        df_alleles_filt.to_csv(self._out_matrix, index_label='ID', sep='\t')
        logger.info(f'Allele matrix exported to: {self._out_matrix}')

        # Calculate the distance matrix
        df_dists = self._calc_distance_matrix(df_alleles_filt)

        # Check the distance matrix
        if df_dists.to_numpy().max() == 0:
            raise ValueError('Empty distance matrix')

        # Export the distance matrix
        df_dists.to_csv(self._out_dists, index_label='ID', sep='\t')
        logger.info(f'Distance matrix exported to: {self._out_dists}')

        # Log the command for GrapeTree
        logger.info(f'You can construct a phylogeny using: grapetree --profile {self._out_matrix} --method MSTreeV2')
