import numpy as np
import matplotlib.pyplot as plt
import os, termcolor, matplotlib, configparser
from pydarm.darm import DARMModel
import scipy.signal as signal
import foton, ipdb

matplotlib.rcParams.update({'font.size': 14})
matplotlib.rcParams.update({'figure.figsize': (14,9)})

def asd_to_rms(f, asd, reverse=False):
    rms_curve = np.zeros_like(asd)
    if not reverse:
        for i in range(len(f)-1, -1, -1):
            rms = np.sqrt(np.sum(asd[i:]**2))
            rms_curve[i] = rms
    else:
        for i in range(0, len(f)):
            rms = np.sqrt(np.sum(asd[:i]**2))
            rms_curve[i] = rms
    print()
    return rms_curve

if __name__ == '__main__':

    freq_50pwm, co2y_ac_50pwm = np.loadtxt(os.path.join(os.path.dirname(__file__),'CO2Y_AC_50PWM.txt'), unpack=True)
    dc_50pwm = 3500 # counts
    freq_cw, co2y_ac_cw = np.loadtxt(os.path.join(os.path.dirname(__file__),'CO2Y_AC_CW.txt'), unpack=True)
    dc_cw = 3880 # counts

    freq, H = signal.freqs_zpk([0], [20], 10**(105.4/20), worN=freq_50pwm*2*np.pi)

    uncalibrated_50pwm = np.abs ( co2y_ac_50pwm / H )
    uncalibrated_cw = np.abs ( co2y_ac_cw / H )

    uncalibrated_dc_50pwm = dc_50pwm / 510
    uncalibrated_dc_cw = dc_cw / 510

    rin_50pwm = uncalibrated_50pwm / uncalibrated_dc_50pwm
    rin_cw = uncalibrated_cw / uncalibrated_dc_cw 

    fig = plt.figure(figsize=(14, 9))
    axs = fig.subplots(1, 1)
    
    axs.loglog(freq_50pwm, rin_50pwm, label='Laser 50% PWM 5kHz', linewidth=.1,
               color = 'C2')
    axs.loglog(freq_50pwm, asd_to_rms(freq_50pwm, rin_50pwm), ls= '--',
               linewidth=.1, color='C2')


    axs.loglog(freq_50pwm, rin_cw, label='Laser on CW', linewidth=.1,
               color='C1')
    axs.loglog(freq_50pwm, asd_to_rms(freq_cw, rin_cw), ls= '--',
               linewidth=.1, color='C1')

    axs.set_xlabel("Frequency [Hz]")
    axs.set_ylabel("RIN (cts/cts)/rtHz")

    axs.grid(True, 'major', 'both', alpha=.5)
    axs.grid(True, 'minor', 'both', alpha=.2)
    axs.legend()

    fig.suptitle("RIN CO2Y Laser OUT : 14 Jan 2025 : H1:IOP-OAF_L0_MADC3_TP_CH10/12")

    fig.savefig(os.path.join(os.path.dirname(__file__), 'figures',
                             '20250114_co2y_rin.pdf'))