import carmcmc as cm
from gatspy.periodic import LombScargleFast
import matplotlib.pyplot as plt
import numpy as np

def csample_from_files(datafile, chainfile, p, q):
    data = np.loadtxt(datafile)

    times, tind = np.unique(data[:,0], return_index=True)
    data = data[tind, :]
    
    chain = np.loadtxt(chainfile)
    assert chain.shape[1] == p + q + 5, 'dimension mismatch'

    return cm.CarmaSample(data[:,0], data[:,1], data[:,2], None, q=q, trace=chain[:,:-2], loglike=chain[:,-2], logpost=chain[:,-1])

def normalised_lombscargle(ts, ys, dys, oversampling=5, nyquist_factor=3):
    model = LombScargleFast().fit(ts, ys, dys)

    pers, pows = model.periodogram_auto(oversampling=oversampling, nyquist_factor=nyquist_factor)
    fs = 1.0/pers

    T = np.max(ts) - np.min(ts)

    mu = 1/T*np.trapz(ys, ts)
    s2 = 1/T*np.trapz(np.square(ys-mu), ts)

    return fs, s2*pows/np.trapz(pows, fs)


def plot_psd_sample_data(sample, oversampling=5, nyquist_factor=3):
    psd_low, psd_high, psd_med, fs = sample.plot_power_spectrum(doShow=False)

    plt.clf()

    plt.loglog(fs, psd_med, '-b', alpha=0.33)
    plt.fill_between(fs, psd_low, psd_high, color='b', alpha=0.17)

    fs, psd = normalised_lombscargle(sample.time, sample.y, sample.ysig, oversampling=oversampling, nyquist_factor=nyquist_factor)

    bw = fs[-1] - fs[0]
    T = sample.time[-1] - sample.time[0]

    s2 = 1/T*np.trapz(np.square(sample.ysig), sample.time)
    noise_level = s2/bw
    levels = noise_level*np.sqrt(sample.get_samples('measerr_scale'))
    plt.axhline(np.median(levels), color='g', alpha=0.33)
    plt.fill_between(fs, np.percentile(levels, 84)+0*fs, np.percentile(levels, 16)+0*fs, color='g', alpha=0.17)

    plt.loglog(fs, psd, '-r', alpha=0.33)

def plot_psd_sample_draw(sample, loc='upper left', oversampling=5, nyquist_factor=3):
    fs, psd = normalised_lombscargle(sample.time, sample.y, sample.ysig, oversampling=oversampling, nyquist_factor=nyquist_factor)

    ys_draw = sample.predict(sample.time, bestfit='random')[0]

    fs, dpsd = normalised_lombscargle(sample.time, ys_draw, sample.ysig, oversampling=oversampling, nyquist_factor=nyquist_factor)

    plt.loglog(fs, psd, '-k', label='Data', alpha=0.5)
    plt.loglog(fs, dpsd, '-b', label='Prediction', alpha=0.5)
    plt.legend(loc=loc)

[(7, 'numpy.loadtxt', 'np.loadtxt', (['datafile'], {}), True, 'import numpy as np\n'), (9, 'numpy.unique', 'np.unique', (['data[:, (0)]'], {'return_index': '(True)'}), True, 'import numpy as np\n'), (12, 'numpy.loadtxt', 'np.loadtxt', (['chainfile'], {}), True, 'import numpy as np\n'), (15, 'carmcmc.CarmaSample', 'cm.CarmaSample', (['data[:, (0)]', 'data[:, (1)]', 'data[:, (2)]', 'None'], {'q': 'q', 'trace': 'chain[:, :-2]', 'loglike': 'chain[:, (-2)]', 'logpost': 'chain[:, (-1)]'}), True, 'import carmcmc as cm\n'), (34, 'matplotlib.pyplot.clf', 'plt.clf', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (36, 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd_med', '"""-b"""'], {'alpha': '(0.33)'}), True, 'import matplotlib.pyplot as plt\n'), (37, 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['fs', 'psd_low', 'psd_high'], {'color': '"""b"""', 'alpha': '(0.17)'}), True, 'import matplotlib.pyplot as plt\n'), (50, 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd', '"""-r"""'], {'alpha': '(0.33)'}), True, 'import matplotlib.pyplot as plt\n'), (59, 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd', '"""-k"""'], {'label': '"""Data"""', 'alpha': '(0.5)'}), True, 'import matplotlib.pyplot as plt\n'), (60, 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'dpsd', '"""-b"""'], {'label': '"""Prediction"""', 'alpha': '(0.5)'}), True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': 'loc'}), True, 'import matplotlib.pyplot as plt\n'), (23, 'numpy.max', 'np.max', (['ts'], {}), True, 'import numpy as np\n'), (23, 'numpy.min', 'np.min', (['ts'], {}), True, 'import numpy as np\n'), (25, 'numpy.trapz', 'np.trapz', (['ys', 'ts'], {}), True, 'import numpy as np\n'), (47, 'numpy.median', 'np.median', (['levels'], {}), True, 'import numpy as np\n'), (18, 'gatspy.periodic.LombScargleFast', 'LombScargleFast', ([], {}), False, 'from gatspy.periodic import LombScargleFast\n'), (26, 'numpy.square', 'np.square', (['(ys - mu)'], {}), True, 'import numpy as np\n'), (28, 'numpy.trapz', 'np.trapz', (['pows', 'fs'], {}), True, 'import numpy as np\n'), (44, 'numpy.square', 'np.square', (['sample.ysig'], {}), True, 'import numpy as np\n'), (48, 'numpy.percentile', 'np.percentile', (['levels', '(84)'], {}), True, 'import numpy as np\n'), (48, 'numpy.percentile', 'np.percentile', (['levels', '(16)'], {}), True, 'import numpy as np\n')]