import numpy as np

import matplotlib.pylab as plt

import physconst

def spectrum(
        T0 = 2.e7,
        dr = 0.02,
        dL = 0.00,
        rp = 0.0,
        lpow = 1,
        rpow = 1,
        nt = 200,
        nv = 100,
        r0 = 1e6,
        r1 = 1e6,
        T1 = 2.5e7,
        fig = None,
        ):
    """
    Code to experiment with resulting RMS of oscillations as a function of energy.

    Parameters
    ----------
    T0 : float
        temperature of mean
    dr : float
        fractional radius variation
    dL : float
        frational luminosity variation
    lpow : integer
        power on L variation, higher makes curve more spiky
    rpow : integer
        power on r variation, higher makes curve more spiky
    nt : integer
        number of points in T direction
    nv : integer
        number of point in frequency (v) direction
    r0 : float
        effective equilibrium radius of varying componnent
    r1 : float
        effective radius of constant componnent
    T1 : float
        temperature of constant contribution
    fig : `~matplotlib.figure.Figure`, optional
        figure to use, generate new one if unspecified.


    """

    h = physconst.PLANCK
    k = physconst.KB
    c = physconst.CLIGHT
    e = physconst.EV
    sb = physconst.SB
    pi = np.pi

    def B(v, T):
        return np.array(2*h*v**3/(c**2*(np.exp(h*v/(k*T))-1)), dtype=np.float64)

    def f(x, pow):
        """
        make asymetric function for light curve
        """
        fx = (1 + np.sin(x))**pow
        return fx - np.average(fx)

    A0 = 4 * pi * r0**2
    L0 = A0 * sb * T0**4
    A1 = 4 * pi * r1**2
    L1 = A1 * sb * T1**4

    v = physconst.EV * 1000 / h * (1 + 9 * np.arange(nv+1) / nv)
    p = np.pi * 2 * np.arange(nt+1) /nt
    E = h * v / (1e3 * e) # E in keV
    r = r0 * (1 + dr * f(p + rp * 2 * np.pi, rpow))
    # dL is change of total L
    L = L0 + dL * f(p, lpow) * (L0 + L1)
    A = 4 * pi * r**2

    F = L / A
    T = (F / sb)**0.25
    b = B(v[np.newaxis, :], T[:, np.newaxis])
    Lv = b * A[:, np.newaxis]

    Lv0 = B(v[np.newaxis, :], T0) * A0
    lv = Lv / Lv0

    # add persistent second componnent
    b1 = B(v[np.newaxis, :], T1)
    Lv1 = b1 * A1
    lv = (Lv + Lv1) / (Lv0 + Lv1)

    s = np.var(lv, axis=0)**0.5
    if fig is None:
        fig = plt.figure(figsize=(10,4))
    else:
        fig.clear()
    ax = fig.add_subplot(121)
    ax.plot(E, s * 100)
    ax.set_xscale('log')
    ax.set_ylabel('Fraction RMS (%)')
    ax.set_xlabel('E (keV)')

    def center(x):
        xc = np.ndarray(len(x)+1)
        xc[1:-1] = 0.5 * (x[1:] + x[:-1])
        xc[0] = 0.5 * (3 * x[0] - x[1])
        xc[-1] = 0.5 * (3 * x[-1] - x[-2])
        return xc

    ax = fig.add_subplot(122)
    cm = ax.pcolormesh(center(E), center(p), (lv-1)*100)
    ax.set_xscale('log')
    ax.set_ylabel('Phase (rad)')
    ax.set_xlabel('E (keV)')
    fig.colorbar(cm, label='relative flux (%)')
    fig.tight_layout()