#!/usr/bin/python3

from math import *
import subprocess
import numpy as np
from scipy.fft import fft, ifft, fftfreq, next_fast_len

def zLoPass(bk, f, r, c):
        return (bk)/(r*(1j*2*pi*f*c)+1)
        #return (bk)*(1/(1j*2*pi*f*c))/(r+(1/(1j*2*pi*f*c)))
        #return (abs(f))*(1/(1j*2*pi*abs(f)*c))/(r+(1/(1j*2*pi*abs(f)*c)))

def get_Zcomp_type2(bk, f, rz, cz, cp):
        #return 7e-5*bk*( (rz*(2*pi*f*cz*1j) + 1.0) / ( (2*pi*f*cp*1j)*(rz*(2*pi*f*cz*1j) + 1.0) + (2*pi*f*cz*1j)) )
        omega = 2*pi*f
        zz = rz + 1.0/(omega*cz*1j)
        return 7e-5*bk / ( (omega*cp*1j) + (1.0/zz) )


def step(a, a0):
        return 1.0*(a>a0)

def ramp(t):
        return np.piecewise(t, [t<0.0, (t>=0.0)&(t<=1.0), t>1.0], [0.0,  lambda t: 0.7 + 4.0*t,  0.0] )

def pwfn_Icaox(t, tpre, ipre, tcur,icur):
        tprev=tpre-tpre[2]
        iprev=ipre
        tcurr=tcur-tcur[0]
        icurr=icur+(ipre[2]-icur[0])
        return np.piecewise(t, \
                [t<tprev[0], (t>=tprev[0])&(t<tprev[1]), (t>=tprev[1])&(t<tcurr[0]), (t>=tcurr[0])&(t<tcurr[1]), (t>=tcurr[1])&(t<tcurr[2]) , t>=tcurr[2]], \
                [0.0, \
                        lambda t: iprev[0]+(t-tprev[0])*((iprev[1]-iprev[0])/(tprev[1]-tprev[0])), \
                        lambda t: iprev[1]+(t-tprev[1])*((icurr[0]-iprev[1])/(tcurr[0]-tprev[1])), \
                        lambda t: icurr[0]+(t-tcurr[0])*((icurr[1]-icurr[0])/(tcurr[1]-tcurr[0])), \
                        lambda t: icurr[1]+(t-tcurr[1])*((icurr[2]-icurr[1])/(tcurr[2]-tcurr[1])), \
                0.0] \
        )


ptsps=int(2400)
npad_pre=10
nper=100
npad_post=9
nper_tot=npad_pre+nper+1+npad_post
tlen=1.0
ttot=tlen*(npad_pre+nper+1+npad_post)
tauscale=1.0/(120)
print("#  next_fast_len:",next_fast_len(ptsps))
tprev = np.linspace(-(nper+npad_pre)*tlen, 0.0, round(tauscale*(nper+npad_pre)*ptsps), endpoint=False)      #prepadding and previous history
tcurr = np.linspace(0.0, tlen, round(tauscale*ptsps), endpoint=False)                                           #current period to solve
tpad = np.linspace(tlen, (npad_post+1)*tlen, round(tauscale*(npad_post)*ptsps), endpoint=False)                 #postpadding
t=np.hstack((tprev, tcurr, tpad))
idx0=-round(tauscale*(npad_post+1)*ptsps)
idxf=-round(tauscale*(npad_post)*ptsps)
if t[idx0]!=0.0 or t[idxf]!=1.0:
        print("check tcurr boundaries?", idx0, t[idx0], idxf, t[idxf])
#ft = np.exp(-0.0*t)*np.cos(2 * np.pi * 5 * t) + 0.5 * np.cos(2 * np.pi * 11 * t ) + 0.25 * np.sin(2 * np.pi * 107 * t +np.pi/4)
#ft = -(2.0/np.pi)+np.exp(-0.1*t)*abs(np.sin(2 * np.pi * 60 * t))
#ft = -(2.0/np.pi)+np.exp(-0.0*t)*(np.sin(2 * np.pi * 60 * t + 1.5*np.sin(2*np.pi*sqrt(2000)*t)))
#ft = sum([1.0*np.sin(2*np.pi*n*t) for n in range(1,250)])
#ft = 0.1*((step(t, 0.2)-step(t, 0.7)-0.0) + (step(t, 0.49)-step(t, 0.99)-0.0))
ta=np.array([0.0, 0.7, 1.0])
tb=np.array([0.0, 0.68, 1.0])
ty=np.array([0.0, 0.9, 1.0])
tz=np.array([0.0, 0.98, 1.0])
ia=np.array([0.35, -0.35, 0.35])
ib=np.array([0.35, -0.34, 0.37])
ic=np.array([0.35, -0.1, 0.65])
iy=np.array([0.45, -0.45, 0.45])
iz=np.array([0.49, -0.49, -0.29])
#ft = 0.005*( pwfn_Icaox(t+1, ty,iy-0.45, ta,ic) + 0*2.3*np.sin(2*np.pi*120*tauscale*t) )
ft = 0.005*( pwfn_Icaox(t+1, ta,ia, tb,ib) + pwfn_Icaox(t, ta,ic, ta,ia) + 0.1*2.3*(np.sin(2*np.pi*120*tauscale*t)) )
ramp0 = ramp(t)
#quit()

rz=2.92e5; cz=3.18e-7; cp=4.54e-8

fw = fft(ft, norm="forward")
freq=fftfreq(round(tauscale*nper_tot*ptsps), d=tauscale*(nper_tot*tlen)/(ptsps))
print("freq[1]={}  freq[{}]={}  tauscale={}    t0=t[{}]={}  tf=t[{}]={}  sum(fw)={:1.6e}".format(freq[1], (ptsps/2)-1, freq[int(ptsps/2)-1:int(ptsps/2)+1], tauscale, idx0, t[idx0], idxf, t[idxf], sum(fw)))
fwlp = np.hstack((0, get_Zcomp_type2(fw[1:], freq[1:], rz, cz, cp)))
ftlp = ifft(fwlp, norm="forward")
dftlp = ifft(fwlp*2.0*np.pi*1j*freq, norm="forward")
print("DC components: fw[0]={:1.6e},  ftlp[0]={:1.6e}".format(fw[0], ftlp[0]))
ftlp+=get_Zcomp_type2(1, 1/(2*pi*1j), 0, cz, cp).real*fw[0]*tauscale*(t-t[idx0])-ftlp[idx0]
print("set: ftlp[{}]={:1.6e}   so ftlp[{}]={:1.6e}".format(idx0, ftlp[idx0], idxf, ftlp[idxf]))
dftlp+=get_Zcomp_type2(1, 1/(2*pi*1j), 0, cz, cp).real*fw[0]*tauscale

gpp=subprocess.Popen(['gnuplot', '-p'], shell=False, universal_newlines=True, stdin=subprocess.PIPE)
gpp.stdin.write("set term qt size 1900,950 \n")
gpp.stdin.write("$ftdat << EOD \n")
for z in range(0, len(fw)):
#for z in range(len(fw)-600, len(fw)):
        ll = [ t[z], tauscale*t[z], freq[z], ramp0[z] ]
        ll.extend([ a
                for b in [ft, ftlp, fw, fwlp, dftlp]
                for a in [b[z].real, b[z].imag, abs(b[z]), np.angle(b[z])]
        ])
        gpp.stdin.write( ("{} "*len(ll)+" \n").format(* ll ) )
gpp.stdin.write("EOD \n")
gpp.stdin.write("plot $ftdat u 1:5 w l lw 2 title 'ft', $ftdat u 1:9 w l lw 2 title'fw' \n")
#gpp.stdin.write("plot $ftdat u 2:5 w l lw 2, $ftdat u 2:9 w l lw 2 \n")
#gpp.stdin.write("plot $ftdat u 3:(log10($15)) w l lw 2, $ftdat u 3:(log10($19)) w l lw 2 \n")
#gpp.stdin.write("plot $ftdat u 3:16 w p lw 2, $ftdat u 3:20 w p lw 2 \n")
gpp.stdin.flush()