#!/usr/bin/python3

import numpy as np
from numpy import pi
from operator import itemgetter
import matplotlib.pyplot as plt
from scipy.optimize import fsolve

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,fc):
    """ izracunam eps in tand iz s parametrov in dolzin """
    s11, s21, s12, s22, f = itemgetter(*s.keys())(s) 
    c = 299792458 #hitrost svetlobe v vakuumu. 

    K = (s11**2 - s21**2 + 1)/2/s11
    G = K + np.sqrt(K**2 - 1)
    Gm = K - np.sqrt(K**2 - 1)
    G[np.abs(G) > 1] = Gm[np.abs(G) > 1]
    T = (s11 + s21 - G)/(1-(s11+s21)*G)
    Lambda2 = (1j/2/pi/L * np.log(T))**2
    l0g2 = ( (f/c)**2 - (fc/c)**2 )**-1
    print(l0g2*Lambda2)
    test_plot(f, np.abs(l0g2*Lambda2))

    return 0

#Nicolson-Ross-Weir metoda -> NUMERICNO NESTABILNA
#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)
#
#
#naumi kako napisat sistem nelinearnih enacb in na kaj jih resevat!
#def iterative(s,fc,L,Lair):
#    s11, s21, s12, s22, f = itemgetter(*s.keys())(s) 
#    c = 299792458 #hitrost svetlobe v vakuumu. 
#
#    p0 = 1j*np.sqrt( (2*np.pi*f/c)**2-(2*np.pi*fc/c)**2)
#    z = lambda p,L: np.exp(-p*L)
#    G = lambda p: (p0-p)/(p0+p) #predpostavim permeabilnost = 1
#
#    f1 = lambda p,L,Lair: np.abs(s21) - np.abs(z(p,L)*(1-G(p)**2)/(1-z(p,L)**2*G(p)**2))
#    f2 = lambda p,L,Lair: np.abs(s11) - np.abs(G(p)*(1-z(p,L)**2)/(1-z(p,L)**2*G(p)**2))
#    f3 = lambda p,L,Lair: (s21*s12-s11*s22)-np.exp(-p0*(Lair-L))*(z(p,L)**2-G(p)**2)/(1-z(p,L)**2*G(p)**2)
#
#    p,l,l_air = fsolve([f1,f2,f3],[p0,L,Lair])
#
#    return p,l,l_air
#

def sistem_enacb(s, fc, beta, L):
    """ sistem enacb za iterativno iskanje dielektricnosti
    predpostavim mu=1 
    
    X = [g, eps, T, G]"""
    s11, s21, s12, s22, f = itemgetter(*s.keys())(s) 
    c = 299792458 #hitrost svetlobe v vakuumu. 

    g0 = 2j*pi*f/c*np.sqrt(1-(fc/f)**2)

    # desne strani enacb
    T = lambda g: np.exp(-g*L) 
    g = lambda eps: 2j*pi*f/c*np.sqrt(eps-(fc/f)**2) # g
    G = lambda g: (g0-g)/(g0+g)

    # pomozne funkcije za zadnjo enacbo
    h1 = lambda T,G: T*(1-G**2)/(1-T**2*G**2)
    h2 = lambda T,G: G*(1-T**2)/(1-T**2*G**2)


    f1 = lambda X: X[2]-T(X[0]) # T
    f2 = lambda X: X[0]-g(X[1]) # g
    f3 = lambda X: X[3]-G(X[0]) #G
    f4 = lambda X: s21+beta*s11-h1(X[2],X[3])-beta*h2(X[2],X[3])

    return lambda x: [f1(x),f2(x),f3(x),f4(x)]

def iterative(s, L, fc):
    c = 299792458 #hitrost svetlobe v vakuumu. 
    f = s['f']
    g0 = 2j*pi*f/c*np.sqrt(1-(fc/f)**2)
    f = sistem_enacb(s, fc, 0.5, L)
    return f


def measured_group_delay(f, P):
    phase = np.unwrap(np.angle(P))
    test_plot(f,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()

### testni klici
a=s2p_to_narray('teflon.s2p')
b = narray_to_s(a)
locals().update(b)
print(iterative(b,6e-2,6.557e9))
#c = s_to_eps(b, 6e-2,6.557e9)