import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import leastsq as lsq
from tqdm import tqdm
import batman

def transit(time, t0, per, rp, a, b, ecc, w, u, law):
    """
    -----------
    Parameters:
    -----------
    time : numpy.ndarray
        time at which the flux is calculated
    t0 : float
        transit central time (in days/or the same as period)
    per : float
        orbital period of exoplanet (in days)
    rp : float
        square-root of tansit depth
    a : float
        scaled semi-major axis
    b : float
        impact parameter
    ecc : float
        eccentricity of the orbit
    w : float
        longitude of peri-astron passage (in deg)
    u : array like
        limb-darkening coefficients
    law : str
        name of the limb-darkening law
        linear, quadratic, nonlinear
    -----------
    return
    -----------
    numpy.ndarray
        array containing transit lightcurve
    """
    para = batman.TransitParams()
    para.t0 = t0
    para.per = per
    para.rp = rp
    para.a = a
    cosi = b/a
    in1 = 180*np.arccos(cosi)/np.pi
    para.inc = in1
    para.ecc = ecc
    para.w = w
    para.u = u
    para.limb_dark = law
    m = batman.TransitModel(para, time)
    fl = m.light_curve(para)
    return fl


def spam(time, per, rp, a, b, u, ecc=0., w=90., t0=0.):
    """
    -----------
    Parameters:
    -----------
    time : numpy.ndarray
        time at which the flux is calculated
    per : float
        orbital period of exoplanet (in days)
    rp : float
        square-root of tansit depth
    a : float
        scaled semi-major axis
    b : float
        impact parameter
    u : array like
        non-linear limb-darkening coefficients
    ecc : float
        eccentricity of the orbit
        default is 0.
    w : float
        longitude of peri-astron passage (in deg)
        default is 90 deg
    t0 : float
        transit central time (in days/or the same as period)
        default is 0.
    -----------
    return
    -----------
    float, float
        SPAM LDCs
    """
    synthetic_flux = transit(time, t0, per, rp, a, b, ecc, w, u, "nonlinear")
    def resid(x):
        model = transit(time, t0, per, rp, a, b, ecc, w, x, "quadratic")
        residuals = synthetic_flux - model
        return residuals
    u1_guess, u2_guess = (12./35.)*u[0] + u[1] + (164./105.)*u[2] + 2.*u[3], (10./21.)*u[0] - (34./63.)*u[2] - u[3]
    soln = lsq(resid, x0 = [u1_guess, u2_guess])
    return soln[0][0], soln[0][1]

def mc_spam(time, per, per_err, rp, rp_err, a, a_err, b, b_err, u, ecc=0., w=90., t0=0.):
    """
    -----------
    Parameters:
    -----------
    time : numpy.ndarray
        time at which the flux is calculated
    per, per_err : float
        orbital period and error in it (in days)
    rp, rp_err : float
        square-root of tansit depth and its error
    a, a_err : float
        scaled semi-major axis and error in it
    b, b_err : float
        impact parameter and its error
    u : array like
        non-linear limb-darkening coefficients
    ecc : float
        eccentricity of the orbit
        default is 0.
    w : float
        longitude of peri-astron passage (in deg)
        default is 90 deg
    t0 : float
        transit central time (in days/or the same as period)
        default is 0.
    -----------
    return
    -----------
    numpy.ndarray, numpy.ndarray
        arrays containing distribution in MC-SPAM LDCs
    """
    period = np.random.normal(per, per_err, 1000)
    rp1 = np.random.normal(rp, rp_err, 1000)
    ar1 = np.random.normal(a, a_err, 1000)
    b1 = np.random.normal(b, b_err, 1000)
    u1_mcs, u2_mcs = np.zeros(1000), np.zeros(1000)
    for i in tqdm(range(len(period))):
        u1_mcs[i], u2_mcs[i] = spam(time, period[i], rp1[i], ar1[i], b1[i], u)
    return u1_mcs, u2_mcs