import numpy as np

from statsmodels.tsa.arima_process import arma_generate_sample
from statsmodels.tsa.arma_mle import Arma


# TODO: still refactoring problem with cov_x
# copied from sandbox.tsa.arima.py
def mcarma22(niter=10, nsample=1000, ar=None, ma=None, sig=0.5):
    """run Monte Carlo for ARMA(2,2)

    DGP parameters currently hard coded
    also sample size `nsample`

    was not a self contained function, used instances from outer scope
      now corrected

    """
    # nsample = 1000
    # ar = [1.0, 0, 0]
    if ar is None:
        ar = [1.0, -0.55, -0.1]
    # ma = [1.0, 0, 0]
    if ma is None:
        ma = [1.0, 0.3, 0.2]
    results = []
    results_bse = []
    for _ in range(niter):
        y2 = arma_generate_sample(ar, ma, nsample + 1000, sig)[-nsample:]
        y2 -= y2.mean()
        arest2 = Arma(y2)
        rhohat2a, cov_x2a, infodict, mesg, ier = arest2.fit((2, 2))
        results.append(rhohat2a)
        err2a = arest2.geterrors(rhohat2a)
        sige2a = np.sqrt(np.dot(err2a, err2a) / nsample)
        # print('sige2a', sige2a,
        # print('cov_x2a.shape', cov_x2a.shape
        # results_bse.append(sige2a * np.sqrt(np.diag(cov_x2a)))
        if cov_x2a is not None:
            results_bse.append(sige2a * np.sqrt(np.diag(cov_x2a)))
        else:
            results_bse.append(np.nan + np.zeros_like(rhohat2a))
    return np.r_[ar[1:], ma[1:]], np.array(results), np.array(results_bse)


def mc_summary(res, rt=None):
    if rt is None:
        rt = np.zeros(res.shape[1])
    nanrows = np.isnan(res).any(1)
    print("fractions of iterations with nans", nanrows.mean())
    res = res[~nanrows]
    print("RMSE")
    print(np.sqrt(((res - rt) ** 2).mean(0)))
    print("mean bias")
    print((res - rt).mean(0))
    print("median bias")
    print(np.median((res - rt), 0))
    print("median bias percent")
    print(np.median((res - rt) / rt * 100, 0))
    print("median absolute error")
    print(np.median(np.abs(res - rt), 0))
    print("positive error fraction")
    print((res > rt).mean(0))


if __name__ == "__main__":

    # short version
    #    true, est, bse = mcarma22(niter=50)
    #    print(true
    #    # print(est
    #    print(est.mean(0)

    """niter 50, sample size=1000, 2 runs
    [-0.55 -0.1   0.3   0.2 ]
    [-0.542401   -0.09904305  0.30840599  0.2052473 ]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.54681176 -0.09742921  0.2996297   0.20624258]


    niter=50, sample size=200, 3 runs
    [-0.55 -0.1   0.3   0.2 ]
    [-0.64669489 -0.01134491  0.19972259  0.20634019]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.53141595 -0.10653234  0.32297968  0.20505973]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.50244588 -0.125455    0.33867488  0.19498214]

    niter=50, sample size=100, 5 runs  --> ar1 too low, ma1 too high
    [-0.55 -0.1   0.3   0.2 ]
    [-0.35715008 -0.23392766  0.48771794  0.21901059]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.3554852  -0.21581914  0.51744748  0.24759245]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.3737861  -0.24665911  0.48031939  0.17274438]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.30015385 -0.27705506  0.56168199  0.21995759]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.35879991 -0.22999604  0.4761953   0.19670835]

    new version, with burnin 1000 in DGP and demean
    [-0.55 -0.1   0.3   0.2 ]
    [-0.56770228 -0.00076025  0.25621825  0.24492449]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.27598305 -0.2312364   0.57599134  0.23582417]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.38059051 -0.17413628  0.45147109  0.20046776]

    [-0.55 -0.1   0.3   0.2 ]
    [-0.47789765 -0.08650743  0.3554441   0.24196087]
    """

    ar = [1.0, -0.55, -0.1]
    ma = [1.0, 0.3, 0.2]
    nsample = 200

    run_mc = True  # False
    if run_mc:
        for sig in [0.1, 0.5, 1.0]:
            import time

            t0 = time.time()
            rt, res_rho, res_bse = mcarma22(niter=100, sig=sig)
            print("\nResults for Monte Carlo")
            print("true")
            print(rt)
            print("nsample =", nsample, "sigma = ", sig)
            print("elapsed time for Monte Carlo", time.time() - t0)
            # 20 seconds for ARMA(2,2), 1000 iterations with 1000 observations
            # sige2a = np.sqrt(np.dot(err2a,err2a)/nsample)
            # print('\nbse of one sample'
            # print(sige2a * np.sqrt(np.diag(cov_x2a))
            print("\nMC of rho versus true")
            mc_summary(res_rho, rt)
            print("\nMC of bse versus zero")  # this implies inf in percent
            mc_summary(res_bse)
            print("\nMC of bse versus std")
            mc_summary(res_bse, res_rho.std(0))
