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, 11, 7, 31), 'numpy.loadtxt', 'np.loadtxt', ({(7, 22, 7, 30): 'datafile'}, {}), '(datafile)', True, 'import numpy as np\n'), ((9, 18, 9, 57), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((12, 12, 12, 33), 'numpy.loadtxt', 'np.loadtxt', ({(12, 23, 12, 32): 'chainfile'}, {}), '(chainfile)', True, 'import numpy as np\n'), ((15, 11, 15, 131), 'carmcmc.CarmaSample', 'cm.CarmaSample', (), '', True, 'import carmcmc as cm\n'), ((34, 4, 34, 13), 'matplotlib.pyplot.clf', 'plt.clf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((36, 4, 36, 45), 'matplotlib.pyplot.loglog', 'plt.loglog', (), '', True, 'import matplotlib.pyplot as plt\n'), ((37, 4, 37, 66), 'matplotlib.pyplot.fill_between', 'plt.fill_between', (), '', True, 'import matplotlib.pyplot as plt\n'), ((50, 4, 50, 41), 'matplotlib.pyplot.loglog', 'plt.loglog', (), '', True, 'import matplotlib.pyplot as plt\n'), ((59, 4, 59, 54), 'matplotlib.pyplot.loglog', 'plt.loglog', (), '', True, 'import matplotlib.pyplot as plt\n'), ((60, 4, 60, 61), 'matplotlib.pyplot.loglog', 'plt.loglog', (), '', True, 'import matplotlib.pyplot as plt\n'), ((61, 4, 61, 23), 'matplotlib.pyplot.legend', 'plt.legend', (), '', True, 'import matplotlib.pyplot as plt\n'), ((23, 8, 23, 18), 'numpy.max', 'np.max', ({(23, 15, 23, 17): 'ts'}, {}), '(ts)', True, 'import numpy as np\n'), ((23, 21, 23, 31), 'numpy.min', 'np.min', ({(23, 28, 23, 30): 'ts'}, {}), '(ts)', True, 'import numpy as np\n'), ((25, 13, 25, 29), 'numpy.trapz', 'np.trapz', ({(25, 22, 25, 24): 'ys', (25, 26, 25, 28): 'ts'}, {}), '(ys, ts)', True, 'import numpy as np\n'), ((47, 16, 47, 33), 'numpy.median', 'np.median', ({(47, 26, 47, 32): 'levels'}, {}), '(levels)', True, 'import numpy as np\n'), ((18, 12, 18, 29), 'gatspy.periodic.LombScargleFast', 'LombScargleFast', ({}, {}), '()', False, 'from gatspy.periodic import LombScargleFast\n'), ((26, 22, 26, 38), 'numpy.square', 'np.square', ({(26, 32, 26, 37): '(ys - mu)'}, {}), '(ys - mu)', True, 'import numpy as np\n'), ((28, 23, 28, 41), 'numpy.trapz', 'np.trapz', ({(28, 32, 28, 36): 'pows', (28, 38, 28, 40): 'fs'}, {}), '(pows, fs)', True, 'import numpy as np\n'), ((44, 22, 44, 44), 'numpy.square', 'np.square', ({(44, 32, 44, 43): 'sample.ysig'}, {}), '(sample.ysig)', True, 'import numpy as np\n'), ((48, 25, 48, 50), 'numpy.percentile', 'np.percentile', ({(48, 39, 48, 45): 'levels', (48, 47, 48, 49): '(84)'}, {}), '(levels, 84)', True, 'import numpy as np\n'), ((48, 57, 48, 82), 'numpy.percentile', 'np.percentile', ({(48, 71, 48, 77): 'levels', (48, 79, 48, 81): '(16)'}, {}), '(levels, 16)', True, 'import numpy as np\n')]