# -*- coding: utf-8 -*-

"""Deterioration algorithm."""

import logging
import math
from itertools import zip_longest
from typing import List, Union

import click
import more_click
import torch

from pykeen.triples import TriplesFactory
from pykeen.typing import TorchRandomHint
from pykeen.utils import ensure_torch_random_state

__all__ = [
    'deteriorate',
]

logger = logging.getLogger(__name__)


def deteriorate(
    reference: TriplesFactory,
    *others: TriplesFactory,
    n: Union[int, float],
    random_state: TorchRandomHint = None,
) -> List[TriplesFactory]:
    """Remove n triples from the reference set.

    TODO: take care that triples aren't removed that are the only ones with any given entity
    """
    if reference.create_inverse_triples:
        raise NotImplementedError

    if isinstance(n, float):
        if n < 0 or 1 <= n:
            raise ValueError
        n = int(n * reference.num_triples)

    generator = ensure_torch_random_state(random_state)
    logger.debug('random state %s', random_state)
    logger.debug('generator %s %s', generator, generator.get_state())
    idx = torch.randperm(reference.num_triples, generator=generator)
    logger.debug('idx %s', idx)
    reference_idx, deteriorated_idx = idx.split(split_size=[reference.num_triples - n, n], dim=0)

    first = reference.clone_and_exchange_triples(
        mapped_triples=reference.mapped_triples[reference_idx],
    )

    # distribute the deteriorated triples across the remaining factories
    didxs = deteriorated_idx.split(math.ceil(n / len(others)), dim=0)
    rest = [
        tf.clone_and_exchange_triples(
            mapped_triples=(
                torch.cat([tf.mapped_triples, reference.mapped_triples[didx]], dim=0)
                if didx is not None else
                tf.mapped_triples  # maybe just give same tf? should it be copied?
            ),
        )
        for didx, tf in zip_longest(didxs, others)
    ]

    return [first, *rest]


@click.command()
@more_click.verbose_option
@click.option('--trials', type=int, default=15, show_default=True)
def _main(trials: int):
    from pykeen.datasets import get_dataset
    from pykeen.constants import PYKEEN_EXPERIMENTS
    import numpy as np
    from tabulate import tabulate
    import pandas as pd
    import seaborn as sns
    import matplotlib.pyplot as plt

    n_comb = trials * (trials - 1) // 2
    logger.info(f'Number of combinations: {trials} n Choose 2 = {n_comb}')

    ns = [0.01, 0.05, 0.1, 0.2, 0.3, 0.4]
    rows = []
    for dataset_name in [
        # 'kinships',
        'nations',
        # 'umls',
        # 'codexsmall',
        # 'wn18',
    ]:
        dataset_rows = []
        reference_dataset = get_dataset(dataset=dataset_name)
        for n in ns:
            similarities = []
            for trial in range(trials):
                deteriorated_dataset = reference_dataset.deteriorate(n=n, random_state=trial)
                sim = reference_dataset.similarity(deteriorated_dataset)
                similarities.append(sim)
                rows.append((dataset_name, n, trial, sim))
            dataset_rows.append((n, np.mean(similarities), np.std(similarities)))
        click.echo(tabulate(dataset_rows, headers=[f'{dataset_name} N', 'Mean', 'Std']))

    df = pd.DataFrame(rows, columns=['name', 'n', 'trial', 'sim'])
    tsv_path = PYKEEN_EXPERIMENTS / 'deteriorating.tsv'
    png_path = PYKEEN_EXPERIMENTS / 'deteriorating.png'
    click.echo(f'writing to {tsv_path}')
    df.to_csv(tsv_path, sep='\t', index=False)
    sns.lineplot(data=df, x="n", y="sim", hue="name")
    plt.savefig(png_path, dpi=300)


if __name__ == '__main__':
    _main()
