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)

[((179, 199), 'numpy.loadtxt', 'np.loadtxt', (['datafile'], {}), '(datafile)\n', (189, 199), True, 'import numpy as np\n'), ((219, 261), 'numpy.unique', 'np.unique', (['data[:, (0)]'], {'return_index': '(True)'}), '(data[:, (0)], return_index=True)\n', (228, 261), True, 'import numpy as np\n'), ((301, 322), 'numpy.loadtxt', 'np.loadtxt', (['chainfile'], {}), '(chainfile)\n', (311, 322), True, 'import numpy as np\n'), ((396, 537), 'carmcmc.CarmaSample', 'cm.CarmaSample', (['data[:, (0)]', 'data[:, (1)]', 'data[:, (2)]', 'None'], {'q': 'q', 'trace': 'chain[:, :-2]', 'loglike': 'chain[:, (-2)]', 'logpost': 'chain[:, (-1)]'}), '(data[:, (0)], data[:, (1)], data[:, (2)], None, q=q, trace=\n    chain[:, :-2], loglike=chain[:, (-2)], logpost=chain[:, (-1)])\n', (410, 537), True, 'import carmcmc as cm\n'), ((1061, 1070), 'matplotlib.pyplot.clf', 'plt.clf', ([], {}), '()\n', (1068, 1070), True, 'import matplotlib.pyplot as plt\n'), ((1076, 1117), 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd_med', '"""-b"""'], {'alpha': '(0.33)'}), "(fs, psd_med, '-b', alpha=0.33)\n", (1086, 1117), True, 'import matplotlib.pyplot as plt\n'), ((1122, 1184), 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['fs', 'psd_low', 'psd_high'], {'color': '"""b"""', 'alpha': '(0.17)'}), "(fs, psd_low, psd_high, color='b', alpha=0.17)\n", (1138, 1184), True, 'import matplotlib.pyplot as plt\n'), ((1712, 1749), 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd', '"""-r"""'], {'alpha': '(0.33)'}), "(fs, psd, '-r', alpha=0.33)\n", (1722, 1749), True, 'import matplotlib.pyplot as plt\n'), ((2169, 2219), 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'psd', '"""-k"""'], {'label': '"""Data"""', 'alpha': '(0.5)'}), "(fs, psd, '-k', label='Data', alpha=0.5)\n", (2179, 2219), True, 'import matplotlib.pyplot as plt\n'), ((2224, 2281), 'matplotlib.pyplot.loglog', 'plt.loglog', (['fs', 'dpsd', '"""-b"""'], {'label': '"""Prediction"""', 'alpha': '(0.5)'}), "(fs, dpsd, '-b', label='Prediction', alpha=0.5)\n", (2234, 2281), True, 'import matplotlib.pyplot as plt\n'), ((2286, 2305), 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': 'loc'}), '(loc=loc)\n', (2296, 2305), True, 'import matplotlib.pyplot as plt\n'), ((766, 776), 'numpy.max', 'np.max', (['ts'], {}), '(ts)\n', (772, 776), True, 'import numpy as np\n'), ((779, 789), 'numpy.min', 'np.min', (['ts'], {}), '(ts)\n', (785, 789), True, 'import numpy as np\n'), ((804, 820), 'numpy.trapz', 'np.trapz', (['ys', 'ts'], {}), '(ys, ts)\n', (812, 820), True, 'import numpy as np\n'), ((1553, 1570), 'numpy.median', 'np.median', (['levels'], {}), '(levels)\n', (1562, 1570), True, 'import numpy as np\n'), ((605, 622), 'gatspy.periodic.LombScargleFast', 'LombScargleFast', ([], {}), '()\n', (620, 622), False, 'from gatspy.periodic import LombScargleFast\n'), ((843, 861), 'numpy.square', 'np.square', (['(ys - mu)'], {}), '(ys - mu)\n', (852, 861), True, 'import numpy as np\n'), ((889, 907), 'numpy.trapz', 'np.trapz', (['pows', 'fs'], {}), '(pows, fs)\n', (897, 907), True, 'import numpy as np\n'), ((1406, 1428), 'numpy.square', 'np.square', (['sample.ysig'], {}), '(sample.ysig)\n', (1415, 1428), True, 'import numpy as np\n'), ((1620, 1645), 'numpy.percentile', 'np.percentile', (['levels', '(84)'], {}), '(levels, 84)\n', (1633, 1645), True, 'import numpy as np\n'), ((1652, 1677), 'numpy.percentile', 'np.percentile', (['levels', '(16)'], {}), '(levels, 16)\n', (1665, 1677), True, 'import numpy as np\n')]