"""Impolements the core algorithm for the nonadditivity analysis."""

import random
from typing import Any

import pandas as pd
from rdkit import Chem
from tqdm.auto import tqdm

from nonadditivity.utils.log import get_logger
from nonadditivity.utils.math import (
    _calculate_mean_and_std_per_compound,
    _calculate_nonadditivity,
    _calculate_theo_quantiles,
)

logger = get_logger()

sign_switch: dict[int, int] = {0: 1, 1: -1}


def top_distance_changes(var_lhs: str, var_rhs: str) -> bool:
    """Check whether a linker exchange is valid.

    Checks whether the topological distance between two or three anchoring points
    is the same if a linker is exchanged.

    Args:
        var_lhs (str): SMILES of compound 1
        var_rhs (str): SMILES of compound 2

    Returns:
        bool: False if distance is the same.
    """
    fragment_a = Chem.MolFromSmarts(var_lhs)  # type: ignore pylint: disable=E1101
    fragment_b = Chem.MolFromSmarts(var_rhs)  # type: ignore pylint: disable=E1101
    atoms_a = [atom.GetSmarts() for atom in fragment_a.GetAtoms()]
    atoms_b = [atom.GetSmarts() for atom in fragment_b.GetAtoms()]
    matrix_a = Chem.GetDistanceMatrix(fragment_a)  # type: ignore pylint: disable=E1101
    matrix_b = Chem.GetDistanceMatrix(fragment_b)  # type: ignore pylint: disable=E1101

    a_idx1 = atoms_a.index("[*:1]")
    a_idx2 = atoms_a.index("[*:2]")
    b_idx1 = atoms_b.index("[*:1]")
    b_idx2 = atoms_b.index("[*:2]")
    if matrix_a[a_idx1, a_idx2] != matrix_b[b_idx1, b_idx2]:
        return True
    if "[*:3]" in var_lhs:
        a_idx3 = atoms_a.index("[*:3]")
        b_idx3 = atoms_b.index("[*:3]")
        if matrix_a[a_idx1, a_idx3] != matrix_b[b_idx1, b_idx3]:
            return True
        if matrix_a[a_idx2, a_idx3] != matrix_b[b_idx2, b_idx3]:
            return True
    return False


def add_neighbor_dictionary_to_df(
    mmp_dataframe: pd.DataFrame,
    per_compound_dataframe: pd.DataFrame,
    no_chiral: bool,
    verbose: bool = True,
) -> pd.DataFrame:
    """Get Shortest transformation from mmpdb output.

    Checks the data generated by mmpdblib index, and adds a column to the
    v giving neighbors and the respective transformation and constant part
    for every entry. All done inplace.

    !We could also run mmpdb with --smallest-transformation-only, but this
    is not faster than doing it here. As of November 2022!

    Args:
        mmp_dataframe (pd.DataFrame): data generated by mmpdblib
        per_compound_dataframe (pd.DataFrame): dataframe to add the column with
        neigbors to.
        no_chiral (bool, optional): True if you want to skip chiral compounds. Defaults
        to False.
        verbose (bool): wheter to display tqdm status bar. Defaults to True.
    """
    neigh_dict = {key: {} for key in per_compound_dataframe.index.to_numpy()}
    const_dict = {key: {} for key in per_compound_dataframe.index.to_numpy()}
    seen = {key: set() for key in per_compound_dataframe.index.to_numpy()}
    lengths = {key: {} for key in per_compound_dataframe.index.to_numpy()}
    for id_lhs, id_rhs, transf, constant in tqdm(
        zip(
            mmp_dataframe["ID_LHS"].values,
            mmp_dataframe["ID_RHS"].values,
            mmp_dataframe["TRANSFORMATION"].values,
            mmp_dataframe["CONSTANT"].values,
        ),
        desc="Analyzing Neighborhoods",
        total=len(mmp_dataframe["ID_LHS"].values),
        ascii="░▒█",
        disable=not verbose,
    ):
        if no_chiral and "@" in transf:
            continue
        var_lhs, var_rhs = transf.split(">>")
        # Skip pair if the transformation has more than one anchoring point
        # and the topological distance changes between those two (no
        # reason to assume additivity then)
        if "[*:2]" in var_lhs and top_distance_changes(
            var_lhs=var_lhs,
            var_rhs=var_rhs,
        ):
            continue
        newlen = len(transf)
        if id_rhs not in seen[id_lhs]:
            seen[id_lhs].add(id_rhs)
            neigh_dict[id_lhs][id_rhs] = transf
            const_dict[id_lhs][id_rhs] = constant
            lengths[id_lhs][id_rhs] = newlen
            continue

        # Add to neighbor dictionary
        if newlen < lengths[id_lhs][id_rhs]:
            neigh_dict[id_lhs][id_rhs] = transf
            const_dict[id_lhs][id_rhs] = constant
            lengths[id_lhs][id_rhs] = newlen

    for key, value in neigh_dict.items():
        for n_key, n_value in value.items():
            neigh_dict[key][n_key] = (n_value, const_dict[key][n_key])
    per_compound_dataframe["Neighbor_dict"] = neigh_dict.values()
    return per_compound_dataframe


def _get_tranformation_pairs(
    neighbordict: dict[str, dict[str, str]],
) -> dict[str, list[tuple[str, ...]]]:
    """Create dict mapping id pair to transformation.

    Creates dictionary that maps a transformation to a list with
    all compound pairs (ID1, ID2) with the relation Transformation(ID1) -> ID2
    only returns transformation that have more than one entry.

    Args:
        neighbordict (dict[str, dict[str, str]]): dictionary with IDs\
            , their neighbors and transformations.

    Returns:
        dict[str, list[tuple[str, ...]]]: {Transformation: [(ID1, ID2),...]}
    """
    transformation_pairs: dict[str, list[tuple[str, ...]]] = {}
    for compound_id, neighbor_dictionary in neighbordict.items():
        for neighbor_id, transformation in neighbor_dictionary.items():
            transformation_pairs[transformation] = transformation_pairs.get(
                transformation,
                [],
            )
            transformation_pairs[transformation].append(
                (compound_id, neighbor_id),
            )
    return transformation_pairs


def _find_circles(
    neighbordict: dict[str, dict[str, str]],
    transformation_pairs: dict[str, list[tuple[str, ...]]],
    verbose: bool = True,
) -> list[list[str]]:
    r"""Find Double transformation cycles.

    Core algorithm of the nonadditivity package.
    Following Name convention is followed:

    compound_1 --------t1--------> compound_2
        |                            |
        t2                           t3
        \/                           \/
    compound_4 --------t4-------> compound_3

    entries are appended to return list as [compound_1, compound_2, compound_3,
    compound_4]

    Args:
        neighbordict (dict[str, dict[str, str]]): dict mapping compound id to
        a dict with its neighbors and the transformations to them.
        transformation_pairs (dict[str, list[tuple[str, ...]]]): dict mapping
        transformation to tuples of comoound ids that can be mutated into each other by
        the transformation.
        verbose (bool): wheter to display tqdm status bar. Defaults to True.

    Returns:
        list[list[int]]: list of circle lists[compound_1, compound_2, compound_3,
        compound_4]
    """
    already_checked_1 = {key: False for key in neighbordict}
    already_checked_2 = {key: False for key in neighbordict}
    circles = []
    for compound_1, compound_1_neighbors in tqdm(
        neighbordict.items(),
        desc="Finding Double Transformation Cycles",
        ascii="░▒█",
        disable=not verbose,
    ):
        # No circle possible if only one neighbor to compound_1!
        if len(compound_1_neighbors.values()) > 1:
            already_checked_2 = already_checked_1.copy()
            for (
                compound_2,
                compound_1_to_compound_2,
            ) in compound_1_neighbors.items():
                compound_2_neighbors = neighbordict[compound_2]
                if already_checked_1[compound_2]:
                    continue
                for compound_4, compound_3 in transformation_pairs[
                    compound_1_to_compound_2
                ]:
                    if any(
                        [
                            compound_4 not in compound_1_neighbors,
                            compound_3 not in compound_2_neighbors,
                            already_checked_2[compound_4],
                            already_checked_1[compound_3],
                            compound_2 == compound_4,
                            compound_1 == compound_3,
                        ],
                    ):
                        continue
                    if (
                        compound_1_neighbors[compound_4]
                        == compound_2_neighbors[compound_3]
                    ):
                        circles.append(
                            [compound_1, compound_2, compound_3, compound_4],
                        )
                already_checked_2[compound_2] = True
        already_checked_1[compound_1] = True
    return circles


def get_circles(
    per_compound_dataframe: pd.DataFrame,
    verbose: bool = True,
) -> tuple[list[list[str]], dict[str, list[tuple[str, ...]]]]:
    """Find all dtcs for the compounds in the per compound dataframe.

    Wrapper for the core algorithm of the nonadditivity package.
    Removes all

    Args:
        per_compound_dataframe (pd.DataFrame): per compound data
        verbose (bool): If true, more is logged.

    Returns:
        tuple[list[list[str]],  dict[str, list[tuple[str, ...]]]]: all
        double transformation circles as well as the corresponding trans
        formations.
    """
    per_compound_dataframe = per_compound_dataframe[
        per_compound_dataframe.Neighbor_dict != {}
    ]
    neighbordict = {
        key: {k: v[0] for k, v in value.items()}
        for key, value in zip(
            per_compound_dataframe.index.values,
            per_compound_dataframe.Neighbor_dict.values,
        )
    }
    transformation_pairs = _get_tranformation_pairs(neighbordict=neighbordict)
    circles = _find_circles(
        neighbordict=neighbordict,
        transformation_pairs=transformation_pairs,
        verbose=verbose,
    )
    return circles, transformation_pairs


def _randomize_circles(
    circles: list[list[str]],
) -> tuple[list[list[str]], list[int]]:
    """Randomize circles.

    Takes a list of circles and randomly chooses starting point for every circle
    and returns new list of then randomized circles.

    Args:
        circles (list[list[str]]): original circles

    Returns:
        tuple[list[list[str]], list[int]]: randomized circles, shift for every
        randomizer circle
    """
    new_circles = []
    min_indices = []
    for circle in circles:
        min_index = random.choice(range(4))
        new_circles.append((circle + circle)[min_index : min_index + 4])
        min_indices.append(min_index)
    return new_circles, min_indices


def _create_circle_ids(
    circles: list[list[str]],
    property_column: str,
) -> list[str]:
    """Create IDS for all circles.

    Args:
        circles (list[list[str]]): list of circles
        property_column (str): name of the property that is investigated.

    Returns:
        list[str]: list with circle ids.
    """
    return [
        "_".join([str(c[0]), str(c[1]), str(c[2]), str(c[3])]) + f"_{property_column}"
        for c in circles
    ]


def _is_same_series(
    per_compound_dataframe: pd.DataFrame,
    circle: list[str],
) -> bool:
    """True if all compounds of a circle are in the same series.

    Args:
        per_compound_dataframe (pd.DataFrame): per compound dataframe
        circle (list[str]): circle to be checked.

    Returns:
        bool: true if all compounds are in the same series.
    """
    return len({per_compound_dataframe.loc[srn, "Series"] for srn in circle}) == 1


def _get_faulty_entries(
    per_compound_dataframe: pd.DataFrame,
    circles: list[list[str]],
    property_column: str,
    include_censored: bool,
) -> set[str]:
    faulty = set()
    censor_column = f"{property_column}_Censors"
    for index, circle in enumerate(circles):
        for compound in circle:
            if per_compound_dataframe.loc[compound, censor_column] == "NA":
                faulty.add(index)
                continue
            if (
                not include_censored
                and per_compound_dataframe.loc[compound, censor_column]
            ):
                faulty.add(index)
    return faulty


def create_c2c_dataframe(c2c_vals: dict[str, list[Any]]) -> pd.DataFrame:
    """Create dataframe mapping circle to compounds.

    Args:
        c2c_vals (dict[str, list[Any]]): values to be added to new dataframe.

    Returns:
        pd.DataFrame: dataframe with c2c vals.
    """
    c2c_dataframe = pd.DataFrame()
    c2c_dataframe["Circle_ID"] = c2c_vals["circle_ids"]
    c2c_dataframe["Compound_ID"] = c2c_vals["compound_ids"]
    c2c_dataframe["SMILES"] = c2c_vals["smiles"]
    c2c_dataframe["Property"] = c2c_vals["properties"]
    c2c_dataframe["NonadditivityPerCompound"] = c2c_vals["nonadditivities"]
    return c2c_dataframe


def create_na_dataframe(
    na_vals: dict[str, list[Any]],
    randomized_circles: list[list[str]],
    per_compound_dataframe: pd.DataFrame,
) -> pd.DataFrame:
    """Create NAA dataframe containing circles and their nonadditivity.

    Args:
        na_vals (dict[str, list[Any]]): nonadditivity values
        randomized_circles (list[list[str]]): circle ids
        per_compound_dataframe (pd.DataFrame): per compound dataframe

    Returns:
        pd.DataFrame: nonadditivity data
    """
    na_dataframe = pd.DataFrame()
    for i in range(4):
        na_dataframe[f"Compound{i+1}"] = [circle[i] for circle in randomized_circles]
    for i in range(4):
        na_dataframe[f"SMILES{i+1}"] = [
            per_compound_dataframe.loc[circle[i], "SMILES"]
            for circle in randomized_circles
        ]
    for i in range(4):
        na_dataframe[f"Prop_Cpd{i+1}"] = na_vals[f"cpdvals{i}"]

    for i, v in zip([1, 2], [1, 3]):
        na_dataframe[f"Transformation{i}"] = [
            per_compound_dataframe.loc[circle[0], "Neighbor_dict"][circle[v]][0]
            for circle in randomized_circles
        ]
    for name in [
        "Property",
        "Series",
        "Nonadditivity",
        "Theo_Quantile",
        "Circle_ID",
    ]:
        na_dataframe[name] = na_vals[name]
    return na_dataframe


def calculate_na_output(
    circles: list[list[str]],
    per_compound_dataframe: pd.DataFrame,
    property_columns: list[str],
    include_censored: bool,
    series_column: str | None = None,
) -> tuple[pd.DataFrame, list[list[str]], list[str], pd.DataFrame]:
    """Create dataframe with nonadditivity circles.

    Creates a dataframe with an entry for every nonadditivity cicle and property.
    also returns randomized circles and their corresponding circle ids.

    Args:
        circles (list[list[str]]): list of nonadditivity circles
        per_compound_dataframe (pd.DataFrame): per compound dataframe
        property_columns (list[str]): name of investigated property columns
        include_censored (bool): Whether to also use censored values.
        series_column (str | None, optional): name of the series column. Defaults to
        None.

    Returns:
        tuple[pd.DataFrame, list[list[str]], list[str]]: nonadditivity dataframe,
        randomized circles, circle_ids
    """
    na_vals: dict[str, list[float] | list[str]] = {
        key: []
        for key in [
            "Circle_ID",
            "Property",
            "Nonadditivity",
        ]
        + [f"cpdvals{i}" for i in range(4)]
    }
    c2c_vals: dict[str, list[float] | list[str]] = {
        key: []
        for key in [
            "circle_ids",
            "compound_ids",
            "smiles",
            "properties",
            "nonadditivities",
        ]
    }

    randomized_circles = []

    for property_column in property_columns:
        indices_to_remove = _get_faulty_entries(
            per_compound_dataframe=per_compound_dataframe,
            circles=circles,
            property_column=property_column,
            include_censored=include_censored,
        )
        valid_circles = [c for i, c in enumerate(circles) if i not in indices_to_remove]
        if len(valid_circles) < 1:
            _update_na_per_compound(
                per_compound_dataframe=per_compound_dataframe,
                circle=[],
                nonadditivity=0,
                property_column=property_column,
                series_column=series_column,
            )
            continue
        rcircles, _ = _randomize_circles(circles=valid_circles)
        randomized_circles += rcircles
        na_vals["Circle_ID"] += _create_circle_ids(
            circles=rcircles,
            property_column=property_column,
        )
        na_vals["Property"] += [property_column for _ in rcircles]
        for i in range(4):
            na_vals[f"cpdvals{i}"] += [
                per_compound_dataframe.loc[circle[i], property_column]
                for circle in rcircles
            ]
        nonadditivity = [
            _calculate_nonadditivity(
                values=[
                    per_compound_dataframe.loc[index, property_column]
                    for index in circle
                ],
            )
            for circle in rcircles
        ]
        na_vals["Nonadditivity"] += nonadditivity

        for circle, non_add in zip(
            rcircles,
            nonadditivity,
        ):
            _update_na_per_compound(
                per_compound_dataframe=per_compound_dataframe,
                circle=circle,
                nonadditivity=non_add,
                property_column=property_column,
                series_column=series_column,
            )
        cmpds = [i for c in rcircles for i in c]

        c2c_vals["compound_ids"] += cmpds
        c2c_vals["smiles"] += [per_compound_dataframe.loc[i, "SMILES"] for i in cmpds]
        c2c_vals["nonadditivities"] += [
            sign_switch[i % 2] * na
            for circle, na in zip(rcircles, nonadditivity)
            for i, _ in enumerate(circle)
        ]

        c2c_vals["properties"] += [property_column for _ in cmpds]

    c2c_vals["circle_ids"] = [i for i in na_vals["Circle_ID"] for _ in range(4)]
    if series_column:
        na_vals["Series"] = [
            (
                per_compound_dataframe.loc[circle[0], "Series"]
                if _is_same_series(
                    per_compound_dataframe=per_compound_dataframe,
                    circle=circle,
                )
                else ""
            )
            for circle in randomized_circles
        ]
    else:
        na_vals["Series"] = ["" for _ in randomized_circles]

    na_vals["Theo_Quantile"] = _calculate_theo_quantiles(
        series_columns=series_column,
        series_ids=na_vals["Series"],  # type: ignore
        nonadditivities=na_vals["Nonadditivity"],
    )

    return (
        create_na_dataframe(
            na_vals=na_vals,
            randomized_circles=randomized_circles,
            per_compound_dataframe=per_compound_dataframe,
        ),
        randomized_circles,
        na_vals["Circle_ID"],
        create_c2c_dataframe(c2c_vals=c2c_vals),
    )


def _update_na_per_compound(
    per_compound_dataframe: pd.DataFrame,
    circle: list[str],
    nonadditivity: float,
    property_column: str,
    series_column: str | None = None,
) -> None:
    """Adds nonadditivit value to per compound dataframe.

    Args:
        per_compound_dataframe (pd.DataFrame): per compound dataframe.
        circle (list[str]): investigated circle
        nonadditivity (float): nonadditivity value for circle
        series_column (str | None, optional): name of series column. Defaults to None
        property_column (str | None, optional): investigated proeprty column name.
        Defaults to None
    """
    column = f"{property_column}_Nonadditivities"
    if column not in per_compound_dataframe.columns:
        per_compound_dataframe[column] = [
            [] for _ in range(len(per_compound_dataframe))
        ]
    for index, compound in enumerate(circle):
        if series_column:
            is_pure = (
                "pure" if _is_same_series(per_compound_dataframe, circle) else "mixed"
            )
            per_compound_dataframe.loc[compound, column].append(  # type: ignore
                (sign_switch[index % 2] * nonadditivity, is_pure),
            )
        else:
            per_compound_dataframe.loc[compound, column].append(  # type: ignore
                (sign_switch[index % 2] * nonadditivity, None),
            )


def add_mean_and_std_na_to_df(
    per_compound_dataframe: pd.DataFrame,
    property_columns: list[str],
    series_column: str | None = None,
) -> None:
    """Add mean, std and count to per compound dataframe.

    Creates mean, std, and count for all nonadditivities per compound in the per
    compound dataframe and adds new columns with mean, std, and num occurence.
    If series column is given, columns for the pure and the mixed values are added.

    Args:
        per_compound_dataframe (pd.DataFrame): per compound dataframe
        property_columns (list[str]): name of the property columns investigated
        series_column (str | None, optional): series column name. Defaults to None.
    """
    for property_column in property_columns:
        per_compound_values = _calculate_mean_and_std_per_compound(
            per_compound_dataframe=per_compound_dataframe,
            property_column=property_column,
            series_column=series_column,
        )
        if series_column:
            per_compound_dataframe[f"{property_column}_Nonadditivity_Pure"] = [
                value[0]
                for value in per_compound_values[0]  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Pure_SD"] = [
                value[1]
                for value in per_compound_values[0]  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Pure_Count"] = [
                value[2]
                for value in per_compound_values[0]  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Mixed"] = [
                value[0]
                for value in per_compound_values[1]  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Mixed_SD"] = [
                value[1]
                for value in per_compound_values[1]  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Mixed_Count"] = [
                value[2]
                for value in per_compound_values[1]  # type: ignore
            ]
        else:
            per_compound_dataframe[f"{property_column}_Nonadditivity"] = [
                value[0]
                for value in per_compound_values  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_SD"] = [
                value[1]
                for value in per_compound_values  # type: ignore
            ]
            per_compound_dataframe[f"{property_column}_Nonadditivity_Count"] = [
                value[2]
                for value in per_compound_values  # type: ignore
            ]


def run_nonadditivity_core(
    mmp_dataframe: pd.DataFrame,
    per_compound_dataframe: pd.DataFrame,
    property_columns: list[str],
    no_chiral: bool,
    include_censored: bool,
    series_column: str | None = None,
    verbose: bool = True,
) -> tuple[pd.DataFrame, list[list[str]], list[str], pd.DataFrame]:
    """Executes the whole nonadditivity core functions.

    takes an mmp_dataframe and a per compound dataframe and finds
    nonaditivity circles, calculates their nonadditivity, and then
    calculats the nonadditivity per compound.

    Args:
        mmp_dataframe (pd.DataFrame): mmp dataframe from mmpdblib
        per_compound_dataframe (pd.DataFrame): per compound dataframe
        property_columns (list[str]): name of the property columns
        no_chiral (bool): If true, transformation containing @ are skipped.
        include_censored (bool): If true, also include censored values.
        series_column (str | None): name of the series column. Defaults to None.
        verbose (bool, optional): If true, more output is logged. Defaults to True.

    Returns:
        tuple[pd.DataFrame, list[list[str]], list[str]]: nonadditivity dataframe,
        randomized circles, circle_ids
    """
    per_compound_dataframe = add_neighbor_dictionary_to_df(
        mmp_dataframe=mmp_dataframe,
        per_compound_dataframe=per_compound_dataframe,
        no_chiral=no_chiral,
        verbose=verbose,
    )
    circles, _ = get_circles(
        per_compound_dataframe=per_compound_dataframe,
        verbose=verbose,
    )
    logger.info("%s Nonadditivity circles were found.", len(circles))
    (
        na_dataframe,
        randomized_circles,
        circle_ids,
        c2c_dataframe,
    ) = calculate_na_output(
        circles=circles,
        per_compound_dataframe=per_compound_dataframe,
        property_columns=property_columns,
        series_column=series_column,
        include_censored=include_censored,
    )
    add_mean_and_std_na_to_df(
        per_compound_dataframe,
        property_columns,
        series_column=series_column,
    )

    return na_dataframe, randomized_circles, circle_ids, c2c_dataframe
