from pre_compile import average_periodT_module
import matplotlib.pyplot as plt
from mpi4py import MPI
import time
from numba import jit
import numpy as np
import scipy.stats as sts
from numba.pycc import CC

# name of compiled module to create:
cc = CC('average_periodT_module')

# name of function in module, with explicit data types required (4byte=32bit ints and floats)


@cc.export('average_periodT', '(i8, i8, f8, f8, f8, f8[:,:], f8[:,:])')
def average_periodT(S, T, z_0, rho, mu, eps_mat, z_mat):
    first_negative_periodT_lst = []

    for s_ind in range(S):
        z_tm1 = z_0
        for t_ind in range(T):
            e_t = eps_mat[t_ind, s_ind]
            z_t = rho * z_tm1 + (1 - rho) * mu + e_t
            z_mat[t_ind, s_ind] = z_t
            if z_t <= 0:
                first_negative_periodT_lst.append(t_ind + 1)
                break
            else:
                z_tm1 = z_t

    averageT = np.array(first_negative_periodT_lst)
    averageT = np.mean(averageT)

    return averageT


cc.compile()


mu = 3.0
sigma = 1.0
z_0 = mu - 3 * sigma


S = 1000
T = int(4160)

np.random.seed(25)
eps_mat = sts.norm.rvs(loc=0, scale=sigma, size=(T, S))
z_mat = np.zeros((T, S))


def sim():
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    t0 = time.time()

    values = np.linspace(-0.95, 0.95, 200).reshape(10, 20)

    best_T_rho = -9999
    max_periodT = -9999

    for rho in values[rank]:
        averageT = average_periodT_module.average_periodT(
            S, T, z_0, rho, mu, eps_mat, z_mat)
        if max_periodT < averageT:
            max_periodT = averageT
            best_T_rho = rho
    rho_t_array = np.array([best_T_rho, max_periodT])

    # Gather rhos to buffer of expected size/dtype on rank 0
    rho_and_periodT_lst = []
    if rank == 0:
        rho_and_periodT_lst = np.empty([size, 2], dtype='float')
    comm.Gather(sendbuf=rho_t_array, recvbuf=rho_and_periodT_lst, root=0)

    if rank == 0:
        optiomal_rho_s_periodT = max(rho_and_periodT_lst[:, 1])

        for i in range(10):
            if rho_and_periodT_lst[i, 1] == optiomal_rho_s_periodT:
                optiomal_rho = rho_and_periodT_lst[i, 0]
                break

        time_elapsed = time.time() - t0

        print(f"Finished in {time_elapsed} seconds")
        print(f"Optimal rho is {optiomal_rho}")
        print(f"Optiomal rho's period T is {optiomal_rho_s_periodT}")

        # Plot
        x = rho_and_periodT_lst[:, 0]
        y = rho_and_periodT_lst[:, 1]
        plt.xlabel('Rho')
        plt.ylabel('Average Period T')
        plt.plot(x, y)
        plt.savefig("plot.png")


if __name__ == '__main__':
    sim()
