import numpy as np
import pandas as pd
from statsmodels.tsa.arima_process import ArmaProcess
import math
from math import pi
from orbit.exceptions import IllegalArgument


def make_trend(
    series_len, method="rw", arma=[0.25, 0.6], rw_loc=0.0, rw_scale=0.1, seed=1
):
    """Module to generate time-series trend with different methods
    Parameters
    ----------
    series_len: int
        Total length of series
    method: str ['arma', 'rw']
        In case of `'rw'`, a simple random walk process will be used. For `'arma'`, we will use `statsmodels.api` to
        simulate a simple ARMA(1, 1) process
    arma: list
        List [arparams, maparams] of size 2 where used for arma(1) generating process
    rw_loc: float
        Location parameter of random walk generated by `np.random.normal()`
    rw_scale: float
        Scale parameter of random walk generated by `np.random.normal()`
    seed: int
        Seed passed into `np.random.default_rng()`
    Returns
    -------
    np.array-llike
        Simulated trend with length equals `series_len`

    Notes
    -----
        1. ARMA process: https://www.statsmodels.org/stable/generated/statsmodels.tsa.arima_process.ArmaProcess.html
    """
    # make trend
    if method == "rw":
        rw = np.random.default_rng(seed).normal(rw_loc, rw_scale, series_len)
        trend = np.cumsum(rw)
    elif method == "arma":
        arparams = np.array([arma[0]])
        maparams = np.array([arma[1]])
        # add zero-lag and negate
        ar = np.r_[1, -arparams]
        # add zero-lag
        ma = np.r_[1, maparams]
        arma_process = ArmaProcess(ar, ma)
        trend = arma_process.generate_sample(series_len)
    else:
        raise IllegalArgument("Invalid trend method.")

    return trend


def make_seasonality(
    series_len, seasonality, method="discrete", order=3, duration=1, scale=0.05, seed=1
):
    """Module to generate time-series seasonality with different methods
    series_len: int
        Total length of series
    seasonality: int
        For example, seasonality=52 would be a weekly series
    method: str ['discrete', 'fourier']
    order: int
        Used when method = 'fourier' ONLY. Fourier series order to generate seasonality.
    duration: int
        Used when method = 'discrete'. For example, seasonality=52 and duration=7 would be a daily
        series with a constant level on each week.
    scale: float
        Scale parameter of seasonality generation by Normal(0, scale). When method = 'discrete', it is directly used.
        when method = 'fourier', it is used to multiply with the sin-cosine wave generated series.
    seed: int
        Seed passed into `np.random.default_rng()`
    Returns
    -------
    np.array-llike
        Simulated seasonliaty with length equals `series_len`

    Notes
    -----
      1. In case of method = 'discrete', only `seasonality - 1` variables will be generated and the rest is
      directly derived to preserve the property of the sum of seasonality equals 1.
      2.  In case of method = 'fourier', see https://otexts.com/fpp2/complexseasonality.html
    """
    if seasonality > 1:
        if method == "fourier":
            t = np.arange(0, series_len)
            out = []
            for i in range(1, order + 1):
                x = 2.0 * i * np.pi * t / seasonality
                out.append(np.sin(x))
                out.append(np.cos(x))
            out = np.column_stack(out)
            b = np.random.default_rng(seed).normal(0, scale, order * 2)
            seas = np.matmul(out, b)
        else:
            # initialization
            seas = []
            iterations = math.ceil(series_len / duration)
            # initialize vector to be repeated
            init_seas = np.zeros(seasonality)
            init_seas[:-1] = np.random.default_rng(seed).normal(
                0, scale, seasonality - 1
            )
            init_seas[seasonality - 1] = -1 * np.sum(init_seas)
            for idx in range(iterations):
                seas += [init_seas[idx % seasonality]] * duration
            seas = np.array(seas[:series_len])
    else:
        seas = np.zeros(series_len)
    return seas


def make_regression(
    series_len,
    coefs,
    loc=0.0,
    scale=0.5,
    cov=None,
    noise_scale=1.0,
    bias=None,
    relevance=1.0,
    sparsity=0.2,
    seed=1,
):
    """Module to generate multiplicative time-series with trend, seasonality and regression components
    Parameters
    ----------
    series_len: int
        Total length of series
    coefs: 1-D array_like
        Values used as regression coefficients
    loc: float
        Location parameter to generate regressor values
    scale: float
        Scale parameter to generate regressor values
    cov: 2-D array_like, of shape (num_of_regressors, num_of_regressors)
        Covariance of regressors; need to be positive definite
    noise_scale:
        Scale parameter in the white noise generation process
    bias: float
        bias or intercept of the regression component; if None, bias is set to zero.
    relevance: float
        0 to 1; smaller value indicates less number of total useful regressors
    sparsity: float
        0 to 1 to control probability (= 1 - sparsity) at time t of a regressor value > 0
    seed: int
        Seed passed into `np.random.default_rng()`

    Returns
    -------
    x: 2-D array like
        Regressors simulated array. Should be with shape (series_len, num_of_regressors);
    y: 1-D array like
        Regression derived by product of X and coefficients plus noise
    coefs: float
        Coefficients modified in the process due to sparsity; if sparsity=0, it should be identical to coefs inputted
        by user
    """

    num_of_regressors = len(coefs)
    if (relevance > 0.0) and (relevance < 1.0):
        num_of_irr_coefs = int(num_of_regressors * (1 - relevance))
        coefs = np.copy(coefs)
        irr_idx = np.random.choice(num_of_regressors, num_of_irr_coefs, replace=False)
        coefs[irr_idx] = 0.0

    if cov is None:
        x = (
            np.random.default_rng(seed)
            .normal(loc, scale, series_len * num_of_regressors)
            .reshape(series_len, -1)
        )
    else:
        x = np.random.default_rng(seed).multivariate_normal(
            np.array([loc] * num_of_regressors, dtype=np.float64), cov, series_len
        )
    # control probability of regression kick-in
    if (sparsity > 0.0) and (sparsity < 1.0):
        z = (
            np.random.default_rng(seed)
            .binomial(1, 1 - sparsity, series_len * num_of_regressors)
            .reshape(series_len, -1)
        )
        x = x * z

    noise = np.random.default_rng(seed).normal(0, noise_scale, series_len)
    # make observed response
    if bias:
        y = bias + np.matmul(x, coefs) + noise
    else:
        y = np.matmul(x, coefs) + noise

    return x, y, coefs


# def sim_stepwise_coef_data(n, RS, p=3, n_jump=2):
#     """
#     Parameters
#     ----------
#     n : int
#         number of obseravtions
#     RS : int
#         seed
#     p : int
#         number of predictors
#     n_jump
#         number of jumps
#     Returns
#     -------
#     np.ndarray
#
#     Notes
#     -----
#     This code is purely experimental
#
#     """
#     np.random.seed(RS)
#
#     # initializing coefficients at zeros, simulate all coefficient values
#     lev = np.cumsum(np.concatenate((np.array([5.0]), np.random.normal(0, 0.01, n - 1))))
#
#     jump_tp = np.round(np.random.uniform(0, n, (p, n_jump)), 0).astype(np.int32)
#     jump_amp = np.random.normal(0, 0.05, (p, n_jump))
#
#     beta = np.zeros((n, p))
#     beta[0, :] = -2
#     for idx in range(p):
#         beta[jump_tp[idx], idx] = jump_amp[idx]
#     beta = np.cumsum(beta, 0)
#     #     beta += np.random.normal(0, 1e-2, (n, p))
#     # strictly positive coefficients
#     beta = np.exp(beta)
#
#     # simulate regressors
#     covariates = np.random.normal(0, 10, (n, p))
#
#     # observation with noise
#     y = lev + (covariates * beta).sum(-1) + 0.3 * np.random.normal(0, 1, n)
#
#     regressor_col = ["x{}".format(pp) for pp in range(1, p + 1)]
#     data = pd.DataFrame(covariates, columns=regressor_col)
#     data["y"] = y
#     data["date"] = pd.date_range(start="1/1/2018", periods=len(y))
#
#     # hack for p = 3
#     data["beta1"] = beta[:, 0]
#     data["beta2"] = beta[:, 1]
#     data["beta3"] = beta[:, 2]
#
#     return data


# def sim_data_grw(n, RS, p=3):
#     """coefficients curve are geometric random walk like
#
#     Notes
#     -----
#     This code is purely experimental
#     """
#     np.random.seed(RS)
#     beta_init = np.concatenate(
#         (
#             np.random.normal(2.0, 1.0, size=(1, 1)),  # leves
#             np.random.normal(-3.0, 0.05, size=(1, p)),  # regression coefficients
#         ),
#         axis=-1,
#     )
#
#     beta_drift = np.random.normal(0.0, 0.05, size=(n - 1, p + 1))  # drift
#
#     beta = np.concatenate((beta_init, beta_drift), axis=0)
#
#     # geometric random walk
#     beta = np.exp(np.cumsum(beta, 0))
#
#     # simulate regressors
#     covar_lev = np.ones((n, 1))
#     covar = np.concatenate((covar_lev, np.random.normal(0, 10.0, (n, p))), axis=1)
#
#     # observation with noise
#     y = (covar * beta).sum(-1) + 0.3 * np.random.normal(0, 1, n)
#
#     regressor_col = ["x{}".format(pp) for pp in range(1, p + 1)]
#     data = pd.DataFrame(covar[:, 1:], columns=regressor_col)
#     beta_col = ["beta{}".format(pp) for pp in range(1, p + 1)]
#     beta_data = pd.DataFrame(beta[:, 1:], columns=beta_col)
#     data = pd.concat([data, beta_data], axis=1)
#
#     data["y"] = y
#     data["date"] = pd.date_range(start="1/1/2018", periods=len(y))
#
#     return data
#
#
# def sim_data_rw(n, RS, p=3):
#     """coefficients curve are random walk like"""
#     np.random.seed(RS)
#
#     # initializing coefficients at zeros, simulate all coefficient values
#     lev = np.cumsum(np.concatenate((np.array([5.0]), np.random.normal(0, 0.01, n - 1))))
#     beta = np.concatenate(
#         [
#             np.random.uniform(0.05, 0.12, size=(1, p)),
#             np.random.normal(0.0, 0.01, size=(n - 1, p)),
#         ],
#         axis=0,
#     )
#     beta = np.cumsum(beta, 0)
#
#     # simulate regressors
#     covariates = np.random.normal(0, 10, (n, p))
#
#     # observation with noise
#     y = lev + (covariates * beta).sum(-1) + 0.3 * np.random.normal(0, 1, n)
#
#     regressor_col = ["x{}".format(pp) for pp in range(1, p + 1)]
#     data = pd.DataFrame(covariates, columns=regressor_col)
#     beta_col = ["beta{}".format(pp) for pp in range(1, p + 1)]
#     beta_data = pd.DataFrame(beta, columns=beta_col)
#     data = pd.concat([data, beta_data], axis=1)
#
#     data["y"] = y
#     data["date"] = pd.date_range(start="1/1/2018", periods=len(y))
#
#     return data
#
#
# def sim_data_seasonal(n, RS):
#     """coefficients curve are sine-cosine like"""
#     np.random.seed(RS)
#     # make the time varying coefs
#     tau = np.arange(1, n + 1) / n
#     data = pd.DataFrame(
#         {
#             "tau": tau,
#             "date": pd.date_range(start="1/1/2018", periods=n),
#             "beta1": 2 * tau,
#             "beta2": 1.01 + np.sin(2 * pi * tau),
#             "beta3": 1.01 + np.sin(4 * pi * (tau - 1 / 8)),
#             "x1": np.random.normal(0, 10, size=n),
#             "x2": np.random.normal(0, 10, size=n),
#             "x3": np.random.normal(0, 10, size=n),
#             "trend": np.cumsum(
#                 np.concatenate((np.array([1]), np.random.normal(0, 0.1, n - 1)))
#             ),
#             "error": np.random.normal(0, 1, size=n),  # stats.t.rvs(30, size=n),#
#         }
#     )
#
#     data["y"] = (
#         data.x1 * data.beta1 + data.x2 * data.beta2 + data.x3 * data.beta3 + data.error
#     )
#     return data
