from collections import defaultdict
import functools
import itertools
import math
from math import ceil
import tempfile
from typing import Dict, Iterator, List, NamedTuple, Tuple

from flask import make_response, Markup, render_template
import matplotlib
import matplotlib.pyplot as plt
import numba
import numpy as np
import pathos

from steady_cell_phenotype._util import (
    batcher,
    BinCounter,
    complete_search_generator,
    get_phased_trajectory,
    HashableNdArray,
    process_model_text,
    random_search_generator,
    )

matplotlib.use("agg")


class TrajectoryStatistics(NamedTuple):
    trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    data_counts: Dict[HashableNdArray, np.ndarray]
    means: Dict[HashableNdArray, np.ndarray]
    variances: Dict[HashableNdArray, np.ndarray]

    @classmethod
    def make(cls, num_variables: int):
        return TrajectoryStatistics(
                trajectory_length_counts=defaultdict(lambda: BinCounter()),
                data_counts=defaultdict(lambda: np.zeros(shape=0, dtype=int)),
                means=defaultdict(lambda: np.zeros(shape=(0, num_variables))),
                variances=defaultdict(lambda: np.zeros(shape=(0, num_variables))),
                )


def batch_trajectory_process(batch, update_fn) -> TrajectoryStatistics:
    num_samples = batch.shape[0]
    num_variables = batch.shape[1]

    trajectory_statistics = TrajectoryStatistics.make(num_variables)
    # unpack
    trajectory_length_counts = trajectory_statistics.trajectory_length_counts
    data_counts = trajectory_statistics.data_counts
    means = trajectory_statistics.means
    scaled_variances = trajectory_statistics.variances

    for batch_idx in range(num_samples):
        trajectory, phase_state = get_phased_trajectory(batch[batch_idx, :], update_fn)
        trajectory_length_counts[phase_state].add(len(trajectory))

        # extract arrays, to reduce dict lookup again
        data_count = data_counts[phase_state]
        mean = means[phase_state]
        scaled_variance = scaled_variances[phase_state]

        old_len: int = data_count.shape[0]
        trajectory_len: int = trajectory.shape[0]

        if trajectory_len == 0:
            continue

        # resize, if necessary
        if old_len < trajectory_len:
            # extend data_count
            new_data_count: np.ndarray = np.zeros(shape=trajectory_len, dtype=int)
            new_data_count[:old_len] = data_count
            data_count: np.ndarray = new_data_count
            data_counts[phase_state] = new_data_count
            # extend means
            new_mean: np.ndarray = np.zeros(shape=(trajectory_len, num_variables))
            new_mean[:old_len, :] = mean
            mean = new_mean
            means[phase_state] = new_mean
            # extend scaled variance
            new_scaled_variance: np.ndarray = np.zeros(
                    shape=(trajectory_len, num_variables)
                    )
            new_scaled_variance[:old_len] = scaled_variance
            scaled_variance = new_scaled_variance
            scaled_variances[phase_state] = new_scaled_variance

        # do it once
        reversed_trajectory = trajectory[::-1, :]

        # Welford's online algorithm for mean and variance calculation. See Knuth Vol 2, pg 232
        data_count[:trajectory_len] += 1
        old_mean = np.array(mean[:trajectory_len])  # copy
        mean[:trajectory_len] += \
            (reversed_trajectory - old_mean) / np.expand_dims(data_count[:trajectory_len], axis=-1)
        scaled_variance[:trajectory_len] += (reversed_trajectory - old_mean) * \
                                            (reversed_trajectory - mean[:trajectory_len])

    # need to create a new tuple, as means and scaled variances can get clobbered.
    return TrajectoryStatistics(
            trajectory_length_counts=trajectory_length_counts,
            data_counts=data_counts,
            means=means,
            variances=scaled_variances,
            )


def reducer(init_stats: TrajectoryStatistics,
            new_stats: TrajectoryStatistics) -> TrajectoryStatistics:
    # unpack
    trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    data_counts: Dict[HashableNdArray, np.ndarray]
    means: Dict[HashableNdArray, np.ndarray]
    variances: Dict[HashableNdArray, np.ndarray]
    trajectory_length_counts, data_counts, means, variances = init_stats

    # unpack
    next_trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    next_data_counts: Dict[HashableNdArray, np.ndarray]
    next_means: Dict[HashableNdArray, np.ndarray]
    next_scaled_variances: Dict[HashableNdArray, np.ndarray]
    (
        next_trajectory_length_counts,
        next_data_counts,
        next_means,
        next_scaled_variances,
        ) = new_stats

    for phase_point in next_trajectory_length_counts.keys():
        num_variables = next_means[phase_point].shape[1]
        if phase_point not in trajectory_length_counts:
            # just copy it
            trajectory_length_counts[phase_point] = next_trajectory_length_counts[phase_point]
            data_counts[phase_point] = next_data_counts[phase_point]
            means[phase_point] = next_means[phase_point]
            # remember to scale!
            variances[phase_point] = next_scaled_variances[
                                         phase_point
                                     ] / np.expand_dims(data_counts[phase_point], axis=-1)
        else:
            # merge
            trajectory_length_counts[phase_point].add(
                    next_trajectory_length_counts[phase_point]
                    )
            new_max_len = trajectory_length_counts[phase_point].max

            # data counts
            old_data_count: np.ndarray = data_counts[phase_point]
            next_data_count: np.ndarray = next_data_counts[phase_point]
            new_data_counts = np.zeros(shape=new_max_len, dtype=int)
            new_data_counts[: old_data_count.shape[0]] = old_data_count
            new_data_counts[: next_data_count.shape[0]] += next_data_count

            data_counts[phase_point] = new_data_counts

            # means
            old_mean: np.ndarray = means[phase_point]
            next_mean: np.ndarray = next_means[phase_point]
            new_mean: np.ndarray = np.zeros(
                    shape=(new_max_len, num_variables), dtype=np.float64
                    )
            max_overlap = min(old_mean.shape[0], next_mean.shape[0])
            # weighted average on the overlap, copy over on the tail
            new_mean[:max_overlap] = (
                                             np.expand_dims(old_data_count[:max_overlap], axis=-1)
                                             * old_mean[:max_overlap]
                                             + np.expand_dims(next_data_count[:max_overlap], axis=-1)
                                             * next_mean[:max_overlap]
                                     ) / np.expand_dims(new_data_counts[:max_overlap], axis=-1)
            if old_mean.shape[0] > next_mean.shape[0]:
                new_mean[max_overlap:, :] = old_mean[max_overlap:, :]
            else:
                new_mean[max_overlap:, :] = next_mean[max_overlap:, :]

            means[phase_point] = new_mean

            # scaled variances
            old_variance: np.ndarray = variances[phase_point]
            next_scaled_variance: np.ndarray = next_scaled_variances[phase_point]
            new_variance: np.ndarray = np.zeros(
                    shape=(new_max_len, num_variables), dtype=np.float64
                    )
            # "weighted average" on the overlap, copy over on the tail
            old_proportion = np.expand_dims(
                    old_data_count[:max_overlap] / new_data_counts[:max_overlap], axis=-1
                    )
            next_proportion = np.expand_dims(
                    next_data_count[:max_overlap] / new_data_counts[:max_overlap], axis=-1
                    )
            new_variance[:max_overlap] = (
                    old_proportion * old_variance[:max_overlap]
                    + next_proportion
                    * (next_scaled_variance[:max_overlap] / np.expand_dims(next_data_count[:max_overlap], axis=-1))
                    + ((old_mean[:max_overlap] - next_mean[:max_overlap]) ** 2)
                    * old_proportion
                    * next_proportion)
            if old_variance.shape[0] > next_scaled_variance.shape[0]:
                new_variance[max_overlap:, :] = old_variance[max_overlap:, :]
            else:
                new_variance[max_overlap:, :] = next_scaled_variance[
                                                max_overlap:, :
                                                ] / np.expand_dims(next_data_count[max_overlap:], axis=-1)

            variances[phase_point] = new_variance

    return TrajectoryStatistics(trajectory_length_counts, data_counts, means, variances)


def compute_cycles(*,
                   model_text: str,
                   knockouts: Dict[str, int],
                   continuous: Dict[str, bool],
                   num_iterations: int,
                   visualize_variables: Dict[str, bool]):
    """

    Parameters
    ----------
    model_text
    knockouts
    continuous
    num_iterations
    visualize_variables

    Returns
    -------

    """
    # turn boolean-valued dictionary to list of variables to visualize
    visualized_variables: Tuple[str] = tuple(
            var for var in visualize_variables if visualize_variables[var]
            )

    # create an update function and equation system
    variables, update_fn, equation_system = process_model_text(
            model_text, knockouts, continuous
            )

    # compile and warm up the jitter
    update_fn = numba.jit(update_fn)
    update_fn(np.zeros(len(variables), dtype=np.int64))

    # associate variable names with their index in vectors
    variable_idx: Dict[str, int] = dict(zip(variables, range(len(variables))))

    constants: List[str] = equation_system.constant_variables()
    constants_vals: Dict[str, int] = {
        const: int(equation_system[const]) for const in constants
        }

    # decide if we will perform a complete state space search or not
    state_space_size = 3 ** (len(variables) - len(constants))
    perform_complete_search = state_space_size <= num_iterations

    state_generator: Iterator[np.ndarray]
    if perform_complete_search:
        state_generator = complete_search_generator(
                variables=variables, constants_vals=constants_vals
                )
        num_iterations = state_space_size
    else:
        state_generator = random_search_generator(
                num_iterations=num_iterations,
                variables=variables,
                constants_vals=constants_vals,
                )

    max_threads = max(1, min(pathos.multiprocessing.cpu_count() - 1, math.floor(state_space_size / 1000)))

    batch_generator = batcher(
            state_generator, variables, batch_size=min(1000, num_iterations // max_threads)
            )

    trajectory_length_counts: Dict[HashableNdArray, BinCounter] = dict()
    data_counts: Dict[HashableNdArray, np.ndarray] = dict()
    means: Dict[HashableNdArray, np.ndarray] = dict()
    variances: Dict[HashableNdArray, np.ndarray] = dict()
    with pathos.multiprocessing.ProcessPool(nodes=max_threads) as pool:
        trajectory_length_counts, data_counts, means, variances = functools.reduce(
                reducer,
                pool.uimap(
                        batch_trajectory_process, batch_generator, itertools.repeat(update_fn)
                        ),
                TrajectoryStatistics(
                        trajectory_length_counts, data_counts, means, variances
                        ),
                )

    phase_points = trajectory_length_counts.keys()

    # recreate limit cycles from phase point keys
    limit_cycles: Dict[HashableNdArray, List] = dict()  # record limit sets with their ordering. i.e. as cycles
    for phase_point in phase_points:
        cycle = list()
        t_state = phase_point
        state = phase_point.array
        # do-while
        cycle.append(t_state)
        state = update_fn(state)
        t_state = HashableNdArray(state)
        while t_state != phase_point:
            cycle.append(t_state)
            state = update_fn(state)
            t_state = HashableNdArray(state)
        # record cycle
        limit_cycles[phase_point] = cycle

    # create images for each variable
    limit_set_stats_images = {phase_point: dict() for phase_point in phase_points}
    with tempfile.TemporaryDirectory() as tmp_dir_name:
        plt.rcParams["svg.fonttype"] = "none"
        for phase_point, var in itertools.product(phase_points, visualized_variables):
            cycle = limit_cycles[phase_point]
            var_idx = variable_idx[var]

            var_means = np.flipud(means[phase_point][:, var_idx])
            var_stdevs = np.flipud(np.sqrt(variances[phase_point][:, var_idx]))

            plt.figure(figsize=(6, 4))

            # means, since phased will fix on single value at end
            plt.plot(list(var_means) + [cycle[0][var_idx]], color="#1f3d87")
            plt.fill_between(
                    range(len(var_means) + 1),
                    list(var_means - var_stdevs) + [cycle[0][var_idx]],
                    list(var_means + var_stdevs) + [cycle[0][var_idx]],
                    color="grey",
                    alpha=0.25,
                    )
            # dashed cycle in `converging` region
            plt.plot(
                    range(len(var_means) - len(cycle), len(var_means) + 1),
                    [cycle[idx][var_idx] for idx in range(len(cycle))]
                    + [cycle[0][var_idx]],
                    color="#F24F00",
                    linestyle="dashed",
                    )
            # solid cycle in `converged` region
            plt.plot(
                    range(len(var_means), len(var_means) + len(cycle) + 1),
                    [cycle[idx][var_idx] for idx in range(len(cycle))]
                    + [cycle[0][var_idx]],
                    color="#F24F00",
                    )

            plt.title(f"Phased Envelope of trajectories for {var}")

            plt.xlabel("Distance to phase-fixing point on limit cycle")
            max_tick_goal = 40
            tick_gap = max(1, ceil((len(var_means) + len(cycle) + 1) / max_tick_goal))
            x_tick_locations = [
                x
                for x in range(0, len(var_means) + len(cycle) + 1)
                if (x - len(var_means)) % tick_gap == 0
                ]
            x_tick_labels = list(map(lambda x: len(var_means) - x, x_tick_locations))
            plt.xticks(x_tick_locations, x_tick_labels, rotation=90)
            plt.xlim([0, len(var_means) + len(cycle)])
            plt.axvline(len(var_means), linestyle="dotted", color="grey")
            plt.axvline(len(var_means) - len(cycle), linestyle="dotted", color="grey")

            plt.ylabel("State")
            plt.yticks([0, 1, 2])
            for y_val in range(3):
                plt.axhline(y=y_val, linestyle="dotted", color="grey", alpha=0.5)
            plt.ylim([0 - 0.1, 2 + 0.1])

            plt.tight_layout()

            image_filename = "dist-" + str(hash(phase_point)) + "-" + var + ".svg"
            plt.savefig(
                    tmp_dir_name + "/" + image_filename, transparent=True, pad_inches=0.0
                    )
            plt.close()
            with open(tmp_dir_name + "/" + image_filename, "r") as image:
                limit_set_stats_images[phase_point][var] = Markup(image.read())

    # print('var images complete')

    # cycle list sorted by frequency
    cycle_list = list(limit_cycles.values())
    cycle_list.sort(
            key=lambda cyc: trajectory_length_counts[cyc[0]].total(), reverse=True
            )

    # tally of number of cycles of each length, sorted by cycle length
    cycle_len_counts = defaultdict(lambda: 0)
    for cycle in cycle_list:
        cycle_len_counts[len(cycle)] += 1
    cycle_len_counts = tuple(sorted(cycle_len_counts.items()))

    length_distribution_images = dict()
    with tempfile.TemporaryDirectory() as tmp_dir_name:
        # create length distribution plots
        plt.rcParams["svg.fonttype"] = "none"
        for phase_point in phase_points:
            length_dist = trajectory_length_counts[phase_point].trimmed_bins()
            plt.figure(figsize=(4, 3))
            plt.bar(x=range(len(length_dist)), height=length_dist, color="#002868")
            plt.title("Distribution of lengths of paths")
            plt.xlabel("Length of path")
            plt.ylabel("Number of states")
            plt.tight_layout()
            image_filename = "dist-" + str(hash(phase_point)) + ".svg"
            plt.savefig(
                    tmp_dir_name + "/" + image_filename, transparent=True, pad_inches=0.0
                    )
            plt.close()
            with open(tmp_dir_name + "/" + image_filename, "r") as image:
                length_distribution_images[phase_point] = Markup(image.read())

    def get_key(cyc):
        return cyc[0]

    def get_count(cyc):
        return trajectory_length_counts[get_key(cyc)].total()

    cycles = [
        {
            "states": cycle,
            "len": len(cycle),
            "count": get_count(cycle),
            "percent": 100 * get_count(cycle) / num_iterations,
            "len-dist-image": length_distribution_images[get_key(cycle)]
            if get_key(cycle) in length_distribution_images
            else "",
            "limit-set-stats-images": limit_set_stats_images[get_key(cycle)]
            if get_key(cycle) in limit_set_stats_images
            else dict(),
            }
        for cycle in cycle_list]

    # respond with the results-of-computation page
    return make_response(
            render_template(
                    "compute-cycles.html",
                    cycles=cycles,
                    variables=variables,
                    cycle_len_counts=cycle_len_counts,
                    complete_results=perform_complete_search,
                    )
            )
