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

"""smc.py

Implementations of a ABC-SMC algorithm for 3DCORE.
"""

import argparse
import datetime
import heliosat
import logging
import multiprocessing
import numpy as np
import os
import pickle
import py3dcore
import sys
import time

from py3dcore._extmath import cholesky
from py3dcore._extnumba import set_random_seed
from py3dcore.params import _numba_calculate_weights_reduce


class ABCSMC(object):
    t_data = []
    b_data = []
    o_data = []
    mask = []

    name = None

    def __init__(self):
        pass

    def add_observation(self, t_data, b_data, o_data):
        """Add magnetic field observation

        Parameters
        ----------
        t_data : np.ndarray
            Time evaluation array.
        b_data : np.ndarray
            Magnetic field array.
        o_data : np.ndarray
            Observer position array.
        """
        self.t_data.extend(t_data)
        self.b_data.extend(b_data)
        self.o_data.extend(o_data)

        _mask = [1] * len(b_data)
        _mask[0] = 0
        _mask[-1] = 0

        self.mask.extend(_mask)

    def init(self, t_launch, model, **kwargs):
        """ABC SMC initialization.

        Parameters
        ----------
        t_launch : datetime.datetime
            Initial CME launch time.
        model : Base3DCOREModel
            3DCORE model class.

        Other Parameters
        ----------------
        seed : int
            Random seed, by default 42.
        set_params : dict
            Dictionary containing parameters to fix to given value.
        """
        logger = logging.getLogger(__name__)
   
        set_params = kwargs.get("set_params", None)

        self.t_launch = t_launch
        self.model = model
        self.parameters = py3dcore.params.Base3DCOREParameters(
            model.default_parameters())
        self.seed = kwargs.get("seed", 42)

        # fix parameters
        if set_params:
            pdict = self.parameters.params_dict

            for spkey, spval in set_params.items():
                for key in pdict:
                    if key == spkey:
                        logger.info("setting \"%s\"=%.3f", key, spval)
                        pdict[key]["distribution"] = "fixed"
                        pdict[key]["fixed_value"] = spval
                    elif spkey.isdigit() and pdict[key]["index"] == int(spkey):
                        logger.info("setting \"%s\"=%.3f", key, spval)
                        pdict[key]["distribution"] = "fixed"
                        pdict[key]["fixed_value"] = spval

            self.parameters._update_arr()

        # setup variables
        self.iter_i = 0

        self.acc_rej_hist = []
        self.eps_hist = []
        self.timer_hist = []

        self.particles = None
        self.weights = None
        self.kernels = None
        self.epses = None

        self.particles_prev = None
        self.weights_prev = None

    def load(self, path):
        logger = logging.getLogger(__name__)

        if os.path.isdir(path):
            files = os.listdir(path)

            if len(files) > 0:
                files.sort()
                path = os.path.join(path, files[-1])
            else:
                raise FileNotFoundError

            logger.info("loading particle file \"%s\"", path)
        elif os.path.exists(path):
            logger.info("loading particle file \"%s\"", path)
        else:
            raise FileNotFoundError

        with open(path, "rb") as fh:
            data = pickle.load(fh)

        self.t_launch = data["t_launch"]

        if data["model"] == "TorusGH3DCOREModel":
            self.model = py3dcore.models.TorusGH3DCOREModel
        elif data["model"] == "ThinTorusGH3DCOREModel":
            self.model = py3dcore.models.ThinTorusGH3DCOREModel
        elif data["model"] == "ThinTorusNC3DCOREModel":
            self.model = py3dcore.models.ThinTorusNC3DCOREModel
        else:
            raise ValueError("model not recognized", data["model"])

        self.parameters = data["parameters"]
        self.seed = data["seed"]

        self.iter_i = len(data["eps_hist"]) - 2

        self.acc_rej_hist = data["acc_rej_hist"]
        self.eps_hist = data["eps_hist"]
        self.timer_hist = data["timer_hist"]

        self.particles = data["particles"]
        self.weights = data["weights"]
        self.kernels = data["kernels"]
        self.epses = data["epses"]
        self.profiles = data["profiles"]

        self.particles_prev = None
        self.weights_prev = None

        self.t_data = data["t_data"]
        self.b_data = data["b_data"]
        self.o_data = data["o_data"]
        self.mask = data["mask"]

        self.name = os.path.basename(path)

    def run(self, iter_end, particles, **kwargs):
        """Run ABC SMC algorithm.

        Parameters
        ----------
        iter_end : int
            Maximum number of iterations.

        particles : int
            Number of particles per iteration, by default 4096.

        Other Parameters
        ----------------
        eps_quantile: float
            Adaptive threshold stepping, by default 0.5.
        jobs : int
            Number of total jobs, by default 8.
        kernel_mode : str
            Transition kernel mode, by default "lcm".
        output : str
            Output folder.
        runs : int
            Number of model runs per worker, by default 2**16.
        sub_iter_max : int
            Maximum number of sub iterations, by default 50.
        workers : int
            Number of parallel workers, by default 8.
        """
        eps_quantile = kwargs.get("eps_quantile", .5)
        jobs = kwargs.get("jobs", 8)
        kernel_mode = kwargs.get("kernel_mode", "lcm")
        output = kwargs.get("output", None)
        runs = kwargs.get("runs", 16)
        sub_iter_max= kwargs.get("sub_iter_max", 50)
        workers = kwargs.get("workers", 8)

        kill_flag = False

        pool = multiprocessing.Pool(processes=workers)

        logger = logging.getLogger(__name__)

        if len(self.eps_hist) == 0:
            eps_0 = py3dcore.abc.rmse([np.zeros((1, 3))] * len(self.b_data), self.b_data)[0]
            self.eps_hist = [eps_0, eps_0 * 0.98]

            logger.info("starting abc algorithm, eps_0 = %0.2fnT", self.eps_hist[-1])

        for iter_i in range(self.iter_i, iter_end):
            logger.info("starting iteration %i", iter_i)

            timer_iter = time.time()

            if iter_i > 0:
                # switch particles/particles_prev
                _tp = self.particles_prev
                self.particles_prev = self.particles

                _tw = self.weights_prev
                self.weights_prev = self.weights

                # decompose particle kernels
                if kernel_mode == "cm":
                    kernels_lower = cholesky(2 * self.kernels)
                elif kernel_mode == "lcm":
                    kernels_lower = cholesky(2 * self.kernels)
            else:
                kernels_lower = None
            
            sub_iter_i = 0
            boost = 0

            rseed = self.seed + 100000 * iter_i
            _results = pool.starmap(abcsmc_worker, [(iter_i, self.model, self.t_launch,
                                                     self.t_data, self.b_data, self.o_data,
                                                     self.mask, self.parameters, self.eps_hist[-1], 
                                                     rseed + i, self.particles_prev, self.weights_prev, 
                                                     kernels_lower, runs, boost, logger)
                                                    for i in range(jobs)])

            total_runs = jobs * int(2**runs)

             # perform additional runs if insufficient particles are collected
            while True:
                tlens = [len(jp[1]) for jp in _results]
                tlen = sum(tlens)

                self.particles = np.zeros((tlen, 14), dtype=np.float32)
                self.epses = np.zeros((tlen, ), dtype=np.float32)
                self.profiles = np.zeros((tlen, len(self.b_data), 3), dtype=np.float32)

                acc_rej = np.array([0, 0, 0])

                for i in range(0, len(_results)):
                    self.particles[sum(tlens[:i]):sum(tlens[:i + 1])] = _results[i][0]
                    self.epses[sum(tlens[:i]):sum(
                        tlens[:i + 1])] = _results[i][1]
                    self.profiles[sum(tlens[:i]):sum(tlens[:i + 1])] = _results[i][2]

                    acc_rej += _results[i][3]

                logger.info("step %i:%i with (%i/%i) particles", iter_i, sub_iter_i, tlen,
                            particles)

                if tlen > particles:
                    break

                # adaptive run boosting
                dr = 19 - runs - boost

                if dr > 0:
                    exp = particles / ((tlen + 1) * (sub_iter_i + 1))

                    if exp > 8 and dr > 3:
                        boost += 3
                    elif exp > 4 and dr > 2:
                        boost += 2
                    elif exp > 2:
                        boost += 1

                rseed = self.seed + 100000 * iter_i + 1000 * (sub_iter_i + 1)
                _results_ext = pool.starmap(abcsmc_worker, [(iter_i, self.model, self.t_launch,
                                                      self.t_data, self.b_data, self.o_data,
                                                      self.mask, self.parameters, 
                                                      self.eps_hist[-1], rseed + i, 
                                                      self.particles_prev, self.weights_prev, kernels_lower,
                                                      runs, boost, logger)
                                                     for i in range(jobs)])
                _results.extend(_results_ext)

                sub_iter_i += 1
                total_runs += jobs * int(2**(runs+boost))

                # kill conditions
                if sub_iter_i == 5 + boost:
                    if tlen * np.floor(sub_iter_max / 5) < particles:
                        logger.warning("expected to exceed maximum number of sub iterations (%i)",
                                       sub_iter_max)
                        logger.warning("aborting")
                        kill_flag = True
                        break

            if kill_flag:
                break

            logger.info("%.2f%% acc, %.2f%% hit, %.2f%% rej", 100 * acc_rej[0] / total_runs,
                        100 * acc_rej[1] / total_runs, 100 * acc_rej[2] / total_runs)

            if tlen > particles:
                self.particles = self.particles[:particles]
                self.epses = self.epses[:particles]

            if iter_i > 0:
                self.weights = np.ones((particles,), dtype=np.float32)
                self.parameters.weight(self.particles, self.particles_prev, self.weights,
                                       self.weights_prev, self.kernels, False)
                self.weights[np.where(self.weights == np.nan)] = 0
            else:
                self.weights = np.ones((particles,), dtype=np.float32) / particles

            # set new eps
            self.eps_hist.append(np.quantile(self.epses, eps_quantile))

            logger.info("setting new eps: %.3f => %.3f", self.eps_hist[-2], self.eps_hist[-1])

            # compute transition kernels
            if kernel_mode == "cm":
                self.kernels = np.cov(self.particles, rowvar=False, aweights=self.weights)

                # due to aweights sometimes very small numbers are generated
                self.kernels[np.where(self.kernels < 1e-14)] = 0
            elif kernel_mode == "lcm":
                kernels_cm = np.cov(self.particles, rowvar=False, aweights=self.weights)

                # due to aweights sometimes very small numbers are generated
                kernels_cm[np.where(kernels_cm < 1e-14)] = 0

                kernels_cm_inv = np.linalg.pinv(kernels_cm)

                logger.info("generating local kernels")
                self.kernels = np.array(pool.starmap(generate_kernels_lcm, [(i, self.particles, kernels_cm_inv) for i in range(particles)]))
            self.acc_rej_hist.append(acc_rej)

            logger.info("step %i done, %i particles, %.2fM runs in %.2f seconds, (total: %s)",
                        iter_i, particles, total_runs / 1e6, time.time() - timer_iter,
                        time.strftime("%Hh %Mm %Ss", time.gmtime(np.sum(self.timer_hist))))

            self.timer_hist.append(time.time() - timer_iter)

            self.iter_i = iter_i

            if output:
                self.save(output)

        pool.close()

    def save(self, path):
        logger = logging.getLogger(__name__)

        if not os.path.exists(path):
            os.makedirs(path)

        if os.path.isdir(path):
            path = os.path.join(path, "{0:02d}".format(self.iter_i))

        data = {
            "particles": self.particles,
            "weights": self.weights,
            "kernels": self.kernels,
            "epses": self.epses,
            "profiles": self.profiles,
            "acc_rej_hist": self.acc_rej_hist,
            "eps_hist": self.eps_hist,
            "timer_hist": self.timer_hist,
            "t_launch": self.t_launch,
            "model": self.model.__name__,
            "parameters": self.parameters,
            "seed": self.seed,
            "t_data": self.t_data,
            "b_data": self.b_data,
            "o_data": self.o_data,
            "mask": self.mask
        }

        logger.info("saving file \"%s\"", path)

        with open(path, "wb") as fh:
            pickle.dump(data, fh)


def abcsmc_worker(iter_i, model, t_launch, t_data, b_data, o_data, mask,  parameters, 
                  eps, seed, particles, weights, kernels_lower, runs, boost, logger):
    model_obj = model(t_launch, int(2**(runs + boost)),
                      parameters=parameters, use_gpu=False)

    if iter_i == 0:
        model_obj.generate_iparams(seed=seed)
    else:
        set_random_seed(seed)
        model_obj.perturb_iparams(particles, weights, kernels_lower)

    profiles = np.array(model_obj.sim_fields(t_data, o_data))

    error = py3dcore.abc.rmse(profiles, b_data, mask=mask)

    accept_mask = error < eps
    rej_mask = np.sum(error == np.inf)

    rej_error = np.sum((error != np.inf) & (error >= eps))
    acc_count = np.sum(accept_mask)

    return model_obj.iparams_arr[accept_mask], error[accept_mask], \
        np.swapaxes(profiles, 0, 1)[accept_mask], \
        np.array([acc_count, rej_error, rej_mask])


def generate_kernels_lcm(i, particles, kernel_cm_inv):
    distances = np.array([_numba_calculate_weights_reduce(particles[i], particles[j], kernel_cm_inv) for j in range(len(particles))])

    cutoff = np.median(distances)

    return np.cov(particles[np.where(distances < cutoff)], rowvar=False)
