import numpy as np
import control as ctrl
import control.matlab as ctrlmatlab
# import slycot# control.minreal requires slycot. requires linux packages "gfortran liblapack-dev" to build when installing through pip3
import scipy.io as spio
import scipy.signal as signal
import sys
import fotontst as foton


##########
#Constants
##########

tolerance = 10**(-4) # tolerance for minimum realisation of a zpk model.
whiteFreq = np.array([1,1]) # 1 Hz (Poles) to remove 1/f^2 rolloff
samplerate = 16384

##########
#Functions
##########
##zpk_minreal strips out poles and zero's that are too close to eachother, within some tolerance.
#Written by Daniel Brown, and gently adapted by VladBoss_speedmeter4eva to compare in Log space,
#replicating Matlab's minreal caluclation ?just about? perfectly.
def zpk_minreal(z,p,k, tol=1e-9):
	"""
	Minimum realisation of a zpk model. Will cancel out zero and poles that are within
	the requested tolerance. This function operates in Log space, so it gives higher
	freqency poles and zeroes less bias.
	"""
	zs = [_ for _ in z if not np.any(abs(np.log10(p) - np.log10(_))< tol)]
	ps = [_ for _ in p if not np.any(abs(np.log10(z) - np.log10(_))< tol)]

	return signal.ZerosPolesGain(np.array(zs, dtype=np.complex128), np.array(ps, dtype=np.complex128), k)

#freq resp function for ZPKs stolen from pyDARM scripts - for getting DC gain (/scr/filtermant.py). Vlad made a small bugfix (needed for python3 only?) in the first functional line
#scipy is lazy and does this by converting to zpk to polynomial and it is simply not viable for long lists of Zeroes and Poles
def freqrespZPK(sys, w):
	"""
	Compute the frequency response of a continuous ZPK filter
	"""
	s = 1j * np.complex128(w)
	numPoles = len(sys.poles)
	g = sys.gain * np.ones(len(w),dtype='complex128')

	sortedPolesIdx = np.argsort(abs(sys.poles))
	sortedZerosIdx = np.argsort(abs(sys.zeros))

	for i in range(0, numPoles):
		if i < len(sys.zeros):
			tmp = s - sys.zeros[sortedZerosIdx[i]]
		else:
			tmp = 1
		g = g * tmp / (s - sys.poles[sortedPolesIdx[i]])

	return g

#########################
#Import StateSpace filter
#########################

quad = spio.loadmat('H1susdata_O3_20191025_ex_zpk.mat')#This zpk made by matlab. Seems to work better than calculating within the control library from the SS
TSTz = quad['TSTz'].T[0]
TSTp = quad['TSTp'].T[0]
TSTk = quad['TSTk'].T[0][0]

#######################################################
#Normalisation coefficient of the original filter at DC
#######################################################

mpN_f_ref = np.abs(freqrespZPK(signal.ZerosPolesGain(TSTz,TSTp,TSTk), [2*np.pi*0.00001])[0])#This is good to the 7th sig fig - at 0.01 Hz its only good to the 3rd sig fig

##################################################################################
#Produce a normalised, discretised, minimum realisation ZPK array to hand to foton
##################################################################################
##Error: control.minreal is just deleting everything or nothing for various of tolerance settings.
##Hence I use Dan Brown's minreal function instead.
##Doing so at the start, we can convert to a polynomial TF because it is simple enough.

sus_zpk = zpk_minreal(np.concatenate((TSTz,-2*np.pi*whiteFreq),axis=0),TSTp,TSTk,tol=tolerance) #1/f^2 whitening filter added in here. Don't worry about changes in the gain here, because this gets normalised anyway.

if (len(sus_zpk.zeros) >= 20 or len(sus_zpk.poles) >= 20):
	print('\ntolerance variable needs to be increased \nYou have {} Zeros and {} Poles. \nYou must have <=20 each, for foton to generate <=10 SOS functions, the hardcoded limit.\n'.format(len(sus_zpk.zeros),len(sus_zpk.poles)))
	sys.exit(1) #Don't really need to halt script here, as foton output will tell you what is wrong (too many SOS) - this is just handy because it tells you why this would go wrong
else:
	print('\nYou have {} Zeros and {} Poles. \nYou should consider altering the tolerance variable \nto have as near to 20 for each as possible.\n'.format(len(sus_zpk.zeros),len(sus_zpk.poles)))

sus_num, sus_den = signal.zpk2tf(sus_zpk.zeros, sus_zpk.poles, sus_zpk.gain)
sus_tf = ctrl.tf(sus_num, sus_den)
susnorm = sus_tf / np.abs(sus_tf.dcgain())
##Unable to discretise cleanly
susnorm_d = ctrlmatlab.c2d(susnorm, 1/samplerate, 'tustin') #gives correct gain for foton ... but the poles and zeroes are a complete mess. TODO why?

##Want to check your normalised TF before pushing it to foton?
# import matplotlib.pyplot as plt
# mag, phase, omega = ctrlmatlab.bode(susnorm, 2*np.pi*np.logspace(-2, 3, 1000))
# plt.show()

###################
#Write out to foton
###################
ff = foton.FilterFile('/ligo/home/vladimir.bossilkov/foton_with_sus/fotoninject/H1CALEX_py.txt')
ff['PCALX_RX_PD'][5].ZPK_set([susnorm.zero(),susnorm.pole(),susnorm_d.dcgain()],plane = 's') #n-plane input always fails. TODO, should be importing susnorm_d for everything, but the zeroes and poles are mangled (plane = 'n' works in this case).
ff['PCALX_RX_PD'][7].ZPK_set([[],[],mpN_f_ref],plane = 'n') #setting plane here is irrelevant - I'm just maintaining old matlab format
ff.write()

