#!/usr/bin/python3

import numpy as np
from operator import itemgetter
import matplotlib.pyplot as plt

def s2p_to_narray(file):
    """ prebere s2p file v narray. STOLPCI --> VRSTICE!!! """
    return np.loadtxt(file, comments=('!','#')).T

def narray_to_s(narray):
    """ prebran s2p narray posortiram po s parametrih.
        pri tem upostevam vrstni red:
        f s11 s21 s12 s22, kjer sta za vsak s parameter
        dve vrstici: realni in imaginarni del """
    keys = ["s{}{}".format(j,i) for i in range(1,3) for j in range(1,3)]
    s = {}
    for i in range(4):
        s[keys[i]] = narray[(2*i+1),:]+1j*narray[(2*i+2),:]
    s['f'] = narray[0,:]
    return s

def s_to_eps(s, L):
    """ izracunam eps in tand iz s parametrov in dolzin """
    s11, s21, s12, s22, f = itemgetter(*s.keys())(s) 

    X = (s11**2 - s21**2 + 1)/2/s11
    G = X + np.sqrt(X**2 - 1)
    Gm = X - np.sqrt(X**2 - 1)
    G[np.abs(G) > 1] = Gm[np.abs(G) > 1]
    P = (s11 + s21 - G)/(1-(s11+s21)*G)
    Lambda2 = - (1/2/np.pi/L * np.log(1/P))**2 #izmisli resitev za korene kompleksnega logaritma
    # argument korena mora biti 2*pi*n, kjer je n=L/lambda_g

    return measured_group_delay(f,P)

def measured_group_delay(f, P):
    phase = np.unwrap(np.angle(P))
    test_plot(f,phase)
    # zgladim fazo s polinomsko aproksimacijo druge stopnje
    #ph = np.polyfit(f,phase,2)
    #print(ph)
    #faza = ph[2]+ph[1]*f+ph[0]*f**2
    #test_plot(f,faza-phase)
    return -np.diff(phase)/np.diff(f)/2/np.pi

def test_plot(x,*args):
    for y in args:
        plt.plot(x,y)
    plt.show()

a=s2p_to_narray('20.s2p')
b = narray_to_s(a)
locals().update(b)
c= s_to_eps(b,6e-2)
test_plot(f[:-1],c)