import json
import logging
import math
import random

from collections import defaultdict
from pathlib import Path

import tensorflow_datasets as tfds

from icgen.dataset_names import DATASETS as _DATASETS


logger = logging.getLogger(__name__)

# NOTE: dataset infos include test data in aggregates for pixel values


def _get_resolution_bracket(dim):
    return math.floor(math.log2(dim)) - 4


def _sample_resolution(mean_dim, min_resolution, max_resolution, max_log_res_deviation):
    # Work in discrete log resolutions
    min_bracket = _get_resolution_bracket(min_resolution)
    max_bracket = _get_resolution_bracket(max_resolution)
    real_bracket = _get_resolution_bracket(mean_dim)
    real_bracket = max(min(real_bracket, max_bracket), min_bracket)

    lower = max(real_bracket - max_log_res_deviation, min_bracket)
    upper = min(real_bracket + max_log_res_deviation, max_bracket)
    resolution_bracket = random.choice(list(range(lower, upper + 1)))
    return 2 ** (4 + resolution_bracket)


def _sample_num_classes(min_classes, max_classes, available_num_classes):
    lower = min(available_num_classes, min_classes)
    upper = min(available_num_classes, max_classes)
    return random.choice(list(range(lower, upper + 1)))


def _sample_classes(available_num_classes, num_classes):
    classes = random.sample(range(available_num_classes), k=num_classes)
    return set(classes)


def _sample_examples(
    max_examples_per_class,
    min_examples_per_class,
    examples_per_class,
    total_examples,
    classes,
):
    num_classes = len(classes)
    upper_test_per_class = (
        min(class_examples for class_examples in examples_per_class) // 2
    )
    upper_test_per_class = min(upper_test_per_class, total_examples // (10 * num_classes))
    lower_test_per_class = 5

    if upper_test_per_class <= lower_test_per_class:
        test_per_class = lower_test_per_class
    else:
        test_per_class = random.choice(
            list(range(lower_test_per_class, upper_test_per_class))
        )

    remaining_examples = total_examples - test_per_class * num_classes
    lower = max(min_examples_per_class * num_classes, remaining_examples // 10)
    upper = min(max_examples_per_class * num_classes, remaining_examples)
    sample_ratio = random.choice(list(range(lower, upper))) / total_examples

    class_to_dev_samples = dict()
    class_to_test_samples = dict()
    for class_id in classes:
        total_class_examples = examples_per_class[class_id]
        total_non_test_class_examples = total_class_examples - test_per_class

        num_class_samples = max(
            min_examples_per_class + test_per_class,
            int(sample_ratio * total_non_test_class_examples + test_per_class),
        )
        num_class_samples = min(num_class_samples, total_class_examples)

        samples = random.sample(range(total_class_examples), k=num_class_samples)

        class_to_test_samples[class_id] = set(samples[:test_per_class])
        class_to_dev_samples[class_id] = set(samples[test_per_class:])
    return class_to_dev_samples, class_to_test_samples


def _dataset_to_augmented_identifier(
    dataset,
    dataset_info,
    min_resolution=16,
    max_resolution=512,
    max_log_res_deviation=1,
    min_classes=2,
    max_classes=100,
    min_examples_per_class=20,
    max_examples_per_class=100_000,
):
    identifier = dict()

    # Dataset
    identifier["dataset"] = dataset

    # Resolution
    identifier["resolution"] = _sample_resolution(
        dataset_info["mean_dim"], min_resolution, max_resolution, max_log_res_deviation
    )

    # Classes
    num_classes = _sample_num_classes(
        min_classes, max_classes, dataset_info["num_classes"]
    )
    identifier["classes"] = _sample_classes(dataset_info["num_classes"], num_classes)

    # Actual samples
    class_to_dev_samples, class_to_test_samples = _sample_examples(
        max_examples_per_class,
        min_examples_per_class,
        dataset_info["examples_per_class"],
        dataset_info["num_examples"],
        identifier["classes"],
    )

    identifier["class_to_dev_samples"] = class_to_dev_samples
    identifier["class_to_test_samples"] = class_to_test_samples
    return identifier


def _dataset_to_identifier(dataset, dataset_info):
    identifier = dict()

    identifier["dataset"] = dataset

    # Use all classes and examples, valid_ratio classes statically
    identifier["classes"] = list(range(dataset_info["num_classes"]))

    num_classes = dataset_info["num_classes"]
    examples_per_class = dataset_info["examples_per_class"]
    total_examples = dataset_info["num_examples"]

    test_per_class = min(class_examples for class_examples in examples_per_class) // 2
    test_per_class = min(test_per_class, total_examples // (num_classes * 10))
    identifier["class_to_test_samples"] = {
        class_: set(range(test_per_class)) for class_ in range(num_classes)
    }
    identifier["class_to_dev_samples"] = {
        class_: set(range(test_per_class, examples_per_class[class_]))
        for class_ in range(num_classes)
    }

    # 1,...5, === 16 | 32 | 64 | 128 | 256 | 512
    mean_dim = dataset_info["mean_dim"]
    identifier["resolution"] = 2 ** (4 + _get_resolution_bracket(mean_dim))

    return identifier


def _identifier_to_data(identifier, dataset):
    label_mapping = {
        sampled_class: index for index, sampled_class in enumerate(identifier["classes"])
    }

    train_split, test_split = [], []
    class_to_count = defaultdict(int)
    for image, label in dataset:  # TODO: possible memory waste
        class_ = int(label)
        if class_ not in identifier["classes"]:
            continue

        image_id = class_to_count[class_]
        class_to_count[class_] += 1

        is_train_image = image_id in identifier["class_to_dev_samples"][class_]
        is_test_image = image_id in identifier["class_to_test_samples"][class_]
        if is_train_image:
            train_split.append((image, label_mapping[class_]))
        elif is_test_image:
            test_split.append((image, label_mapping[class_]))

    return train_split, test_split


def _load_dataset_info(dataset_key, info_path):
    with Path(info_path, f"{dataset_key}.json").open() as info_file:
        info = json.load(info_file)
    return info


def _load_dataset(dataset_key, data_path, info_path, download=False):
    download_config = tfds.download.DownloadConfig(
        extract_dir=data_path / "_extract",
        download_mode=tfds.download.GenerateMode.REUSE_DATASET_IF_EXISTS,
    )
    download_dir = data_path / "_download"
    download_and_prepare_kwargs = dict(
        download_config=download_config, download_dir=download_dir
    )

    dataset = tfds.load(
        dataset_key,
        split=None,  # return as dict: split_name -> valid_ratio
        data_dir=data_path,
        batch_size=None,  # return full datasets as tensors
        in_memory=False,
        shuffle_files=False,
        download=download,
        as_supervised=True,
        download_and_prepare_kwargs=download_and_prepare_kwargs,
    )
    splits = []
    for _, split in tfds.as_numpy(dataset).items():  # dict: split_name -> valid_ratio
        splits += list(split)
    info = _load_dataset_info(dataset_key, info_path)
    return splits, info


class ICDatasetGenerator:
    def __init__(
        self,
        data_path,
        min_resolution=16,
        max_resolution=512,
        max_log_res_deviation=1,
        min_classes=2,
        max_classes=100,
        min_examples_per_class=20,
        max_examples_per_class=100_000,
    ):
        self._data_path = Path(data_path)
        self._info_path = Path(__file__).parent / "infos"

        self.min_resolution = min_resolution
        self.max_resolution = max_resolution
        self.max_log_res_deviation = max_log_res_deviation
        self.min_classes = min_classes
        self.max_classes = max_classes
        self.min_examples_per_class = min_examples_per_class
        self.max_examples_per_class = max_examples_per_class

        self._dataset_to_info = {
            dataset: _load_dataset_info(dataset, self._info_path) for dataset in _DATASETS
        }

    def __repr__(self):
        return (
            f"ICGen("
            f"data_path={self._data_path}, "
            f"min_resolution={self.min_resolution}, "
            f"max_resolution={self.max_resolution}, "
            f"max_log_res_deviation={self.max_log_res_deviation}, "
            f"min_classes={self.min_classes}, "
            f"max_classes={self.max_classes}, "
            f"min_examples_per_class={self.min_examples_per_class}, "
            f"max_examples_per_class={self.max_examples_per_class})"
        )

    def identifier_to_dataset(self, identifier, download=False):
        # Identifier includes:
        # - dataset
        # - classes
        # - class_to_dev_samples
        # - class_to_test_samples
        logger.info("Constructing data from identifier")
        dataset, dataset_info = _load_dataset(
            identifier["dataset"], self._data_path, self._info_path, download
        )
        dev_data, test_data = _identifier_to_data(identifier, dataset)
        return dev_data, test_data, dataset_info

    def get_identifier(self, dataset=None, datasets=None, augment=False):

        if datasets is None and dataset is None:
            raise ValueError("Need to specify either datasets or dataset")

        if datasets and dataset:
            raise ValueError("Do not specify datasets and dataset")

        if datasets:
            dataset = random.choice(datasets)
        dataset_info = self._dataset_to_info[dataset]
        logger.info(f"Sampling identifier for dataset {dataset}")

        if augment:
            return _dataset_to_augmented_identifier(dataset, dataset_info)
        else:
            return _dataset_to_identifier(dataset, dataset_info)

    def get_dataset(self, dataset=None, datasets=None, augment=False, download=False):
        identifier = self.get_identifier(
            dataset=dataset, datasets=datasets, augment=augment
        )
        return self.identifier_to_dataset(identifier, download=download)
