from typing import Sequence
import statsmodels.api as sm
import numpy as np
import pandas as pd

from ..types import ValueStd
from ..utils import trim_zeros
from .. import formulas


def growth_factor(ys: Sequence) -> ValueStd:
    """
    Infer the  growth factor from series of data points ys.

    The growth factor "r" is the parameter that best adjust the sequence:

        y[i + 1] = r * y[i].

    Args:
        ys:
            A sequence of data points approximately generated by the growth
            process shown in the description above.

    Returns:
        A ValueStd named tuple of (value, std), with the associated expected
        value and standard deviation.
    """
    ys = np.asarray(ys)
    eta = ys[1:] / ys[:-1]
    N = len(eta)
    mean_eta = eta.mean()
    delta = np.log(mean_eta) - np.log(eta).mean()

    mean_r = (mean_eta + 2 * delta / N) / (1 + 2 * delta / N)
    denom = N / 2 - 2 * delta
    if denom > 0:
        std_r = np.sqrt(delta / denom) * mean_r
    else:
        std_r = float("inf")
    return ValueStd(mean_r, std_r)


def growth_factors(data):
    """
    Compute growth factors for each column of ys in data frame.

    Return a data frame with ["value", "std"] columns with the original columns
    in the index.
    """
    growth_factors = {}
    for key, col in data.items():
        col = trim_zeros(col)
        growth_factors[key] = growth_factor(col)

    return pd.DataFrame(growth_factors, index=["value", "std"]).T


def exponential_extrapolation(ys: Sequence, n: int, append=False) -> np.ndarray:
    """
    Receive a sequence  and return the next n points of the series
    extrapolating from the input data.

    Args:
        ys:
            Input data.
        n:
            Number of points of the extrapolation.
        append:
            If True, returns a concatenation of the input series with the
            extrapolation.
    """
    ys = np.asarray(ys)
    r, dr = growth_factor(ys)
    K = np.log(r)
    extrapolation = ys[-1] * np.exp(K * np.arange(1, n + 1))
    if append:
        return np.concatenate([ys, extrapolation])
    return extrapolation


def R0_from_cases(model, cases, params, method="OLS") -> ValueStd:
    """
    Read curve of cases and adjust the model R0 from cases.
    """
    # Methods that infer the growth ratio between successive observations
    if method.startswith("ratio-"):
        r, dr = growth_ratio_from_cases(cases, method=method[6:])

        R0 = formulas.R0_from_K(model, params, K=np.log(r))
        R0_plus = formulas.R0_from_K(model, params, K=np.log(r - min(dr, 0.9 * r)))
        R0_minus = formulas.R0_from_K(model, params, K=np.log(r + dr))

    # Methods that infer the exponential growth factor
    elif method in ("OLS",):
        K, dK = growth_factor_from_cases(cases, method=method)

        R0 = formulas.R0_from_K(model, params, K=K)
        R0_plus = formulas.R0_from_K(model, params, K=K - dK)
        R0_minus = formulas.R0_from_K(model, params, K=K + dK)

    else:
        raise ValueError(f"invalid method: {method!r}")

    dR0 = abs(R0_plus - R0_minus) / 2
    return ValueStd(R0, dR0)


def growth_ratio_from_cases(curves, method="GGBayes", **kwargs) -> ValueStd:
    """
    Return the growth rate combining the "cases" and "deaths" columns of an
    epidemic curve.

    Args:
        curves:
            A DataFrame with "cases" and "deaths" columns.
        method:
            Statistical method used to infer growth ratio.
            * GGBayes - A Bayesian Gamma distributed Growth rate.

    Keyword Args:
        Additional keyword arguments are passed to the smoothed_diff function.

    See Also:
        :func:`pydemic.fitting.smoothed_diff`
    """

    if method == "GGBayes":
        fn = lambda col: clean_exponential(curves[col], diff=True, **kwargs)
        cases, deaths = map(fn, curves.columns)
        ratios = [growth_factor(cases), growth_factor(deaths)]
        return average_inference(ratios)
    else:
        raise ValueError


def growth_factor_from_cases(curves, method="OLS", **kwargs) -> ValueStd:
    """
    Return the growth rate combining the "cases" and "deaths" columns of an
    epidemic curve.

    Args:
        curves:
            A DataFrame with "cases" and "deaths" columns.
        method:
            Statistical method used to infer the growth factor.
            * 'OLS' - ordinary least squares.

    Keyword Args:
        Additional keyword arguments are passed to the smoothed_diff function.

    See Also:
        :func:`pydemic.fitting.smoothed_diff`
    """

    if method == "OLS":

        def stats(col):
            data = curves[col]
            data = clean_exponential(data, diff=True, **kwargs)
            Y = np.log(data / data[0])
            X = np.arange(len(Y))
            ols = sm.OLS(Y, sm.add_constant(X))
            res = ols.fit()
            _, K = res.params
            ci = res.conf_int()
            dK = (ci[1, 1] - ci[1, 0]) / 2
            return ValueStd(K, dK)

        cases, deaths = map(stats, curves.columns)
        return average_inference([cases, deaths])
    else:
        raise ValueError
