"""Submodule to handle the model parameterization and randomization."""
import logging
from functools import partial
from os import listdir, path  # pylint: disable=no-name-in-module

import yaml

from ..numerical_libs import sync_numerical_libs, xp
from ..util import distributions, dotdict
from ..util.distributions import generic_distribution


def calc_Te(Tg, Ts, n, f):
    """Calculate the latent period."""
    num = 2.0 * n * f / (n + 1.0) * Tg - Ts
    den = 2.0 * n * f / (n + 1.0) - 1.0
    return num / den


def calc_Reff(m, n, Tg, Te, r):
    """Calculate the effective reproductive number."""
    num = 2.0 * n * r / (n + 1.0) * (Tg - Te) * (1.0 + r * Te / m) ** m
    den = 1.0 - (1.0 + 2.0 * r / (n + 1.0) * (Tg - Te)) ** (-n)
    return num / den


def calc_Ti(Te, Tg, n):
    """Calcuate the infectious period."""
    return (Tg - Te) * 2.0 * n / (n + 1.0)


def calc_beta(Te):
    """Derive beta from Te."""
    return 1.0 / Te


def calc_gamma(Ti):
    """Derive gamma from Ti."""
    return 1.0 / Ti


def CI_to_std(CI):
    """Convert a 95% confidence interval to an equivilent stddev (assuming its normal)."""
    lower, upper = CI
    std95 = xp.sqrt(1.0 / 0.05)
    return (upper + lower) / 2.0, (upper - lower) / std95 / 2.0


# TODO move to util
def recursive_dict_update(d, u):
    """Recursive update() for nested dicts."""
    for k, v in u.items():
        if isinstance(v, dict):
            d[k] = recursive_dict_update(d.get(k, {}), v)
        else:
            d[k] = v
    return d


class buckyParams:
    """Class holding all the model parameters defined in the par file, also used to reroll them for each MC run."""

    @sync_numerical_libs
    def __init__(self, par_file=None):
        """Initialize the class, sync up the libs with the parent context and load the par file."""

        # TODO add flags that are read from the yaml
        self.base_params = dotdict({})
        self.param_funcs = dotdict({})
        self.consts = dotdict({})

        self.opt_params = dotdict({})

        if par_file is not None:
            self.update_params_from_file(par_file)

    def update_params_from_file(self, par_file):
        """Update parameter distributions and consts from new yaml file."""
        update_dict = self.read_yml(par_file)
        self.update_params(update_dict)

    def update_params(self, update_dict):
        """Update parameter distributions and consts from nested dicts."""
        self.consts = recursive_dict_update(self.consts, {k: xp.array(v) for k, v in update_dict["consts"].items()})
        if "opt" in update_dict:
            self.opt_params = recursive_dict_update(self.opt_params, update_dict["opt"])
        self.base_params = recursive_dict_update(self.base_params, update_dict)
        self._generate_param_funcs()

    @staticmethod
    def read_yml(par_file):
        """Read in the YAML par file."""
        # TODO check file exists
        # If par_file is a directory, read files in alphanumeric order
        if path.isdir(par_file):
            root = par_file
            files = listdir(par_file)
            files.sort()
        else:
            root = ""
            files = [par_file]
        # Read in first parameter file
        with open(path.join(root, files[0]), "rb") as f:
            d = yaml.safe_load(f)  # nosec
        # Update dictionary with additional parameter files
        for filename in files[1:]:
            with open(path.join(root, filename), "rb") as f:
                dp = yaml.safe_load(f)
                d = recursive_dict_update(d, dp)
        return d

    def generate_params(self):
        """Generate a new set of params by rerolling, adding the derived params and rejecting invalid sets."""
        while True:  # WTB python do-while...
            params = self.reroll_params()
            if self.consts.Te_min < params.Te < params.Tg and params.Ti > self.consts.Ti_min:
                params.consts = self.consts
                return params
            # logging.debug("Rejected params: " + pformat(params))

    def _generate_param_funcs(self):
        """Generate all the partial functions to roll values of the params."""
        # Default standard deviation (as percentage of mean)
        var = self.consts.reroll_variance if "reroll_variance" in self.consts else 0.2
        # Existing parameter functions
        param_funcs = self.param_funcs

        for p, params in self.base_params.items():
            params = params.copy()
            # Collect distribution name
            if "dist" not in params:
                continue
            dist = params.pop("dist")

            # Set default scale if none is specified (could find better method for this)
            if dist == "truncnorm" and "scale" not in params:
                params["scale"] = xp.abs(var * xp.array(params["loc"]))

            # Clip function
            clip = None
            if "clip" in params:
                clip = partial(xp.clip, **params.pop("clip"))
            # Need to clip after interpolation
            elif "age_bins" in params:
                a_min = params.get("a_min")
                a_max = params.get("a_max")
                if a_min is not None or a_max is not None:
                    clip = partial(xp.clip, a_min=a_min, a_max=a_max)

            # Interpolate function
            interp = None
            if "age_bins" in params:
                standard_age_bins = xp.array(self.consts.age_bins)
                age_bins = xp.array(params.pop("age_bins"))
                interp = partial(self.age_interp, x_bins_new=standard_age_bins, x_bins=age_bins)

            # Main distribution function
            if hasattr(distributions, dist):
                base_func = getattr(distributions, dist)
            elif hasattr(xp.random, dist):
                base_func = getattr(xp.random, dist)
            else:
                logging.error("Distribution {} does not exist!".format(dist))
                continue

            params = {k: xp.array(v) for k, v in params.items()}

            param_funcs[p] = partial(generic_distribution, base_func=base_func, params=params, interp=interp, clip=clip)

    def reroll_params(self):
        """Sample each parameter from distribution and calculate derived parameters."""
        return self.calc_derived_params(dotdict({p: f() for p, f in self.param_funcs.items()}))

    @staticmethod
    def age_interp(x_bins_new, x_bins, y):
        """Interpolate parameters define in age groups to a new set of age groups."""
        # TODO we should probably account for population for the 65+ type bins...
        x_bins_new = xp.array(x_bins_new)
        x_bins = xp.array(x_bins)
        if (x_bins_new.shape != x_bins.shape) or xp.any(x_bins_new != x_bins):
            x_mean_new = xp.mean(x_bins_new, axis=1)
            x_mean = xp.mean(x_bins, axis=1)
            return xp.interp(x_mean_new, x_mean, y)
        return y

    def calc_derived_params(self, params):
        """Add the derived params that are calculated from the rerolled ones."""
        En = self.consts.En
        Im = self.consts.Im
        params["Te"] = calc_Te(
            params["Tg"],
            params["Ts"],
            En,
            params["frac_trans_before_sym"],
        )
        params["Ti"] = calc_Ti(params["Te"], params["Tg"], En)
        r = xp.log(2.0) / params["D"]
        params["R0"] = calc_Reff(
            Im,
            En,
            params["Tg"],
            params["Te"],
            r,
        )

        params["SIGMA"] = 1.0 / params["Te"]
        params["GAMMA"] = 1.0 / params["Ti"]
        params["BETA"] = params["R0"] * params["GAMMA"]
        params["SYM_FRAC"] = 1.0 - params["ASYM_FRAC"]
        params["THETA"] = 1.0 / params["H_TIME"]
        params["GAMMA_H"] = 1.0 / params["I_TO_H_TIME"]
        return params

"""Submodule to handle the model parameterization and randomization."""
import logging
from functools import partial
from os import listdir, path  # pylint: disable=no-name-in-module

import yaml

from ..numerical_libs import sync_numerical_libs, xp
from ..util import distributions, dotdict
from ..util.distributions import generic_distribution


def calc_Te(Tg, Ts, n, f):
    """Calculate the latent period."""
    num = 2.0 * n * f / (n + 1.0) * Tg - Ts
    den = 2.0 * n * f / (n + 1.0) - 1.0
    return num / den


def calc_Reff(m, n, Tg, Te, r):
    """Calculate the effective reproductive number."""
    num = 2.0 * n * r / (n + 1.0) * (Tg - Te) * (1.0 + r * Te / m) ** m
    den = 1.0 - (1.0 + 2.0 * r / (n + 1.0) * (Tg - Te)) ** (-n)
    return num / den


def calc_Ti(Te, Tg, n):
    """Calcuate the infectious period."""
    return (Tg - Te) * 2.0 * n / (n + 1.0)


def calc_beta(Te):
    """Derive beta from Te."""
    return 1.0 / Te


def calc_gamma(Ti):
    """Derive gamma from Ti."""
    return 1.0 / Ti


def CI_to_std(CI):
    """Convert a 95% confidence interval to an equivilent stddev (assuming its normal)."""
    lower, upper = CI
    std95 = xp.sqrt(1.0 / 0.05)
    return (upper + lower) / 2.0, (upper - lower) / std95 / 2.0


# TODO move to util
def recursive_dict_update(d, u):
    """Recursive update() for nested dicts."""
    for k, v in u.items():
        if isinstance(v, dict):
            d[k] = recursive_dict_update(d.get(k, {}), v)
        else:
            d[k] = v
    return d


class buckyParams:
    """Class holding all the model parameters defined in the par file, also used to reroll them for each MC run."""

    @sync_numerical_libs
    def __init__(self, par_file=None):
        """Initialize the class, sync up the libs with the parent context and load the par file."""

        # TODO add flags that are read from the yaml
        self.base_params = dotdict({})
        self.param_funcs = dotdict({})
        self.consts = dotdict({})

        self.opt_params = dotdict({})

        if par_file is not None:
            self.update_params_from_file(par_file)

    def update_params_from_file(self, par_file):
        """Update parameter distributions and consts from new yaml file."""
        update_dict = self.read_yml(par_file)
        self.update_params(update_dict)

    def update_params(self, update_dict):
        """Update parameter distributions and consts from nested dicts."""
        self.consts = recursive_dict_update(self.consts, {k: xp.array(v) for k, v in update_dict["consts"].items()})
        if "opt" in update_dict:
            self.opt_params = recursive_dict_update(self.opt_params, update_dict["opt"])
        self.base_params = recursive_dict_update(self.base_params, update_dict)
        self._generate_param_funcs()

    @staticmethod
    def read_yml(par_file):
        """Read in the YAML par file."""
        # TODO check file exists
        # If par_file is a directory, read files in alphanumeric order
        if path.isdir(par_file):
            root = par_file
            files = listdir(par_file)
            files.sort()
        else:
            root = ""
            files = [par_file]
        # Read in first parameter file
        with open(path.join(root, files[0]), "rb") as f:
            d = yaml.safe_load(f)  # nosec
        # Update dictionary with additional parameter files
        for filename in files[1:]:
            with open(path.join(root, filename), "rb") as f:
                dp = yaml.safe_load(f)
                d = recursive_dict_update(d, dp)
        return d

    def generate_params(self):
        """Generate a new set of params by rerolling, adding the derived params and rejecting invalid sets."""
        while True:  # WTB python do-while...
            params = self.reroll_params()
            if self.consts.Te_min < params.Te < params.Tg and params.Ti > self.consts.Ti_min:
                params.consts = self.consts
                return params
            # logging.debug("Rejected params: " + pformat(params))

    def _generate_param_funcs(self):
        """Generate all the partial functions to roll values of the params."""
        # Default standard deviation (as percentage of mean)
        var = self.consts.reroll_variance if "reroll_variance" in self.consts else 0.2
        # Existing parameter functions
        param_funcs = self.param_funcs

        for p, params in self.base_params.items():
            params = params.copy()
            # Collect distribution name
            if "dist" not in params:
                continue
            dist = params.pop("dist")

            # Set default scale if none is specified (could find better method for this)
            if dist == "truncnorm" and "scale" not in params:
                params["scale"] = xp.abs(var * xp.array(params["loc"]))

            # Clip function
            clip = None
            if "clip" in params:
                clip = partial(xp.clip, **params.pop("clip"))
            # Need to clip after interpolation
            elif "age_bins" in params:
                a_min = params.get("a_min")
                a_max = params.get("a_max")
                if a_min is not None or a_max is not None:
                    clip = partial(xp.clip, a_min=a_min, a_max=a_max)

            # Interpolate function
            interp = None
            if "age_bins" in params:
                standard_age_bins = xp.array(self.consts.age_bins)
                age_bins = xp.array(params.pop("age_bins"))
                interp = partial(self.age_interp, x_bins_new=standard_age_bins, x_bins=age_bins)

            # Main distribution function
            if hasattr(distributions, dist):
                base_func = getattr(distributions, dist)
            elif hasattr(xp.random, dist):
                base_func = getattr(xp.random, dist)
            else:
                logging.error("Distribution {} does not exist!".format(dist))
                continue

            params = {k: xp.array(v) for k, v in params.items()}

            param_funcs[p] = partial(generic_distribution, base_func=base_func, params=params, interp=interp, clip=clip)

    def reroll_params(self):
        """Sample each parameter from distribution and calculate derived parameters."""
        return self.calc_derived_params(dotdict({p: f() for p, f in self.param_funcs.items()}))

    @staticmethod
    def age_interp(x_bins_new, x_bins, y):
        """Interpolate parameters define in age groups to a new set of age groups."""
        # TODO we should probably account for population for the 65+ type bins...
        x_bins_new = xp.array(x_bins_new)
        x_bins = xp.array(x_bins)
        if (x_bins_new.shape != x_bins.shape) or xp.any(x_bins_new != x_bins):
            x_mean_new = xp.mean(x_bins_new, axis=1)
            x_mean = xp.mean(x_bins, axis=1)
            return xp.interp(x_mean_new, x_mean, y)
        return y

    def calc_derived_params(self, params):
        """Add the derived params that are calculated from the rerolled ones."""
        En = self.consts.En
        Im = self.consts.Im
        params["Te"] = calc_Te(
            params["Tg"],
            params["Ts"],
            En,
            params["frac_trans_before_sym"],
        )
        params["Ti"] = calc_Ti(params["Te"], params["Tg"], En)
        r = xp.log(2.0) / params["D"]
        params["R0"] = calc_Reff(
            Im,
            En,
            params["Tg"],
            params["Te"],
            r,
        )

        params["SIGMA"] = 1.0 / params["Te"]
        params["GAMMA"] = 1.0 / params["Ti"]
        params["BETA"] = params["R0"] * params["GAMMA"]
        params["SYM_FRAC"] = 1.0 - params["ASYM_FRAC"]
        params["THETA"] = 1.0 / params["H_TIME"]
        params["GAMMA_H"] = 1.0 / params["I_TO_H_TIME"]
        return params

Submodule to handle the model parameterization and randomization. # pylint: disable=no-name-in-module Calculate the latent period. Calculate the effective reproductive number. Calcuate the infectious period. Derive beta from Te. Derive gamma from Ti. Convert a 95% confidence interval to an equivilent stddev (assuming its normal). # TODO move to util Recursive update() for nested dicts. Class holding all the model parameters defined in the par file, also used to reroll them for each MC run. Initialize the class, sync up the libs with the parent context and load the par file. # TODO add flags that are read from the yaml Update parameter distributions and consts from new yaml file. Update parameter distributions and consts from nested dicts. Read in the YAML par file. # TODO check file exists # If par_file is a directory, read files in alphanumeric order # Read in first parameter file # nosec # Update dictionary with additional parameter files Generate a new set of params by rerolling, adding the derived params and rejecting invalid sets. # WTB python do-while... # logging.debug("Rejected params: " + pformat(params)) Generate all the partial functions to roll values of the params. # Default standard deviation (as percentage of mean) # Existing parameter functions # Collect distribution name # Set default scale if none is specified (could find better method for this) # Clip function # Need to clip after interpolation # Interpolate function # Main distribution function Sample each parameter from distribution and calculate derived parameters. Interpolate parameters define in age groups to a new set of age groups. # TODO we should probably account for population for the 65+ type bins... Add the derived params that are calculated from the rerolled ones.