#!/usr/bin/python3

from fenics import *
import numpy as np
import matplotlib.pyplot as plt
from sys import argv

from scipy.constants import epsilon_0, mu_0

primeri = ('koax','izmaknjen','potlacen','kvadratni','stripline')

if argv[1] in primeri:
    problem = argv[1]
else:
    print('Neznan primer. Nadaljujem s koax.')
    problem = 'koax'
shrani = True #spravi vse rezultate v fajle?

mreza = Mesh('geometrija/{}.xml'.format(problem))
domena = MeshFunction('size_t', mreza,
        'geometrija/{}_physical_region.xml'.format(problem))
rob = MeshFunction('size_t', mreza,
        'geometrija/{}_facet_region.xml'.format(problem))

dx = Measure('dx', domain=mreza, subdomain_data=domena)
ds = Measure('ds', domain=mreza, subdomain_data=rob)

V = FunctionSpace(mreza, 'P', 2)

napetost = 100
zila = DirichletBC(V, Constant(napetost), rob, 2)
oklop = DirichletBC(V, Constant(0), rob, 3)
r_pogoji = [zila, oklop]
eps_r = 2.55

u = TrialFunction(V)
v = TestFunction(V)
a = dot(grad(u), grad(v))*dx
L = Constant(0)*v*dx
u = Function(V)
solve(a == L, u, r_pogoji)

plot(mreza, linewidth=0.5)
p = plot(u)
cp = plt.colorbar(p)
cp.set_label('Elektricni potencial V (V)')
if shrani:
    plt.savefig('slike/{}-V.pdf'.format(problem),bbox_inches='tight')
plt.show()

E = -grad(u)
plot(mreza, linewidth=0.5)
p = plot(E, zorder=3)
cp = plt.colorbar(p)
cp.set_label('Elektricno polje E (V/m)')
if shrani:
    plt.savefig('slike/{}-E.pdf'.format(problem),bbox_inches='tight')
plt.show()

Z_w = np.sqrt(mu_0/epsilon_0/eps_r)
cross_z = lambda w: as_vector((-w[1],+w[0]))
H = cross_z(E)/Z_w

plot(mreza, linewidth=0.5)
p = plot(H, zorder=3)
cp = plt.colorbar(p)
cp.set_label('Magnetno polje H (A/m)')
if shrani:
    plt.savefig('slike/{}-H.pdf'.format(problem),bbox_inches='tight')
plt.show()

nhat = FacetNormal(mreza)
cross_t = lambda w,x: w[0]*x[1]-w[1]*x[0]
k = cross_t(nhat, H)
tok = assemble(-k*ds(2))
Z0 = napetost/tok

C = epsilon_0*eps_r/napetost**2 * assemble(E**2 * dx(1))
L = mu_0/tok**2 * assemble(H**2 * dx(1))

print(tok, Z0, C, L)
if shrani:
    with open('izracuni/{}-out.txt'.format(problem), 'w') as out:
        out.write('$I$,{},A\n'.format(tok))
        out.write('$Z_0$,{},\\ohm\n'.format(Z0))
        out.write('$C$,{},pF\\per\\metre\n'.format(C*1e12))
        out.write('$L$,{},nH\\per\\metre'.format(L*1e9))

# nadaljni del kode relevanten samo za koaksilani vod
# izračuna analitske vrednosti in jih za primer polj
# primerja z analitsko resitvijo
if problem == 'koax':
    b = 3
    a = 0.5
    lg = np.log(b/a)
    u_anal = lambda r: napetost/lg*np.log(r/a)
    E_anal = lambda r: napetost/lg/r
    H_anal = lambda r: E_anal(r)/Z_w
    tok_anal = 2*np.pi*napetost/lg/Z_w
    Z0_anal = 1/2/np.pi*Z_w*lg
    C_anal = 2*np.pi*eps_r*epsilon_0/lg
    L_anal = mu_0/2/np.pi*lg

    print(tok_anal, tok)
    print(Z0_anal, Z0)
    print(C_anal, C)
    print(L_anal, L)
    if shrani:
        with open('izracuni/{}-out.txt'.format(problem), 'w') as out:
            out.write('velicina,FEM,analično,enote,napaka\n')
            out.write('$I$,{},{},A,'.format(tok,tok_anal))
            out.write('{}\n'.format((tok-tok_anal)/tok_anal*100))
            out.write('$Z_0$,{},{},\\ohm,'.format(Z0,Z0_anal))
            out.write('{}\n'.format((Z0-Z0_anal)/Z0_anal*100))
            out.write('$C$,{},{},pF\\per\\metre,'.format(C*1e12,C_anal*1e12))
            out.write('{}\n'.format((C-C_anal)/C_anal*100))
            out.write('$L$,{},{},nH\\per\\metre,'.format(L*1e9,L_anal*1e9))
            out.write('{}'.format((L-L_anal)/L_anal*100))

    x = np.linspace(a, b, 100)
    y = np.zeros(len(x))

    primer_ana_e = E_anal(x)
    primer_fem_e = np.array(list(map(project(E),zip(x,y))))[:,0]
    plt.plot(x, primer_ana_e, label='Analitsko')
    plt.plot(x, primer_fem_e, '--', label='FEniCS')
    plt.xlabel('Polmer (m)')
    plt.ylabel('Elektricno polje E (V/m)')
    plt.grid()
    plt.legend()
    if shrani:
        plt.savefig('slike/{}-anal-e.pdf'.format(problem),bbox_inches='tight')
    plt.show()
    
    plt.plot(x, (primer_ana_e-primer_fem_e)/primer_ana_e*100)
    plt.xlabel('Polmer (m)')
    plt.ylabel('Relativna napaka %')
    plt.grid()
    if shrani:
        plt.savefig('slike/{}-napaka-e.pdf'.format(problem),bbox_inches='tight')
    plt.show()

    primer_ana_h = H_anal(x)
    primer_fem_h = np.array(list(map(project(H),zip(x,y))))[:,1]
    plt.plot(x, primer_ana_h, label='Analitsko')
    plt.plot(x, primer_fem_h, '--', label='FEniCS')
    plt.xlabel('Polmer (m)')
    plt.ylabel('Magnetno polje H (V/m)')
    plt.grid()
    plt.legend()
    if shrani:
        plt.savefig('slike/{}-anal-h.pdf'.format(problem),bbox_inches='tight')
    plt.show()
    
    plt.plot(x, (primer_ana_h-primer_fem_h)/primer_ana_h*100)
    plt.xlabel('Polmer (m)')
    plt.ylabel('Relativna napaka %')
    plt.grid()
    if shrani:
        plt.savefig('slike/{}-napaka-h.pdf'.format(problem),bbox_inches='tight')
    plt.show()

if problem == 'stripline':
    def We(W, b):
        if W/b > 0.35:
            return W
        else:
            We_b = W/b-(0.35-W/b)**2
            return We_b*b

    W = 0.8e-3
    b = 1.6e-3
    tmp = b/(We(W,b)+0.441*b)
    Z0strip = 30*np.pi/np.sqrt(eps_r)*tmp
    print(Z0strip)
    print(Z0)
    napaka = (Z0strip-Z0)/Z0*100
    print(napaka)