# -*- coding: utf-8 -*-
#"""
#Created on Fri Jan  9 14:26:54 2015

#@author: lbrieda
#"""
import pylab as pl

EPS0 = 8.85418782e-12
Q = 1.602e-19

node_offset = 0

class Piece:
    def __init__(self,length,nodes,eps_r):
        global node_offset        
        self.length = length
        self.nodes = nodes
        self.eps = eps_r*EPS0
        self.offset = node_offset
        self.dh = length/(nodes-1)
        node_offset = node_offset+nodes-1

#define pieces
pieces = []
pieces.append(Piece(0.05,8,1));
pieces.append(Piece(0.02,6,3));
pieces.append(Piece(0.03,4,1));

#total number of nodes
nn = 1
for p in pieces:
    nn = nn+(p.nodes-1)

#boundaries
phi_bc = (0,10)

#initial potential array
phi = [0]*nn
phi[0] = phi_bc[0]
phi[nn-1] = phi_bc[1]

#charge density   
rho_f = [2e10*Q]*nn  

#solver loop
for it in range(1,1000):
    #iterate over pieces
    for k in range(0,len(pieces)):
        p = pieces[k]
        #set left node per interface b.c. except for first piece
        if (k>0):
            eps_plus = p.eps
            eps_minus= pieces[k-1].eps
            dh_plus = p.dh
            dh_minus = pieces[k-1].dh
            i = p.offset
            sigma_f = rho_f[i]*0.5*(dh_plus+dh_minus)
            phi[i] = 1.0/(dh_minus*eps_plus+dh_plus*eps_minus)*(sigma_f*dh_plus*dh_minus + \
                        eps_plus*dh_minus*phi[i+1] + eps_minus*dh_plus*phi[i-1])
            
        #on piece-internal nodes, solve poisson's equation nabla^2 phi = -rho/eps
        for i in range(p.offset+1,p.offset+p.nodes-1):
            if (i<nn-1):
                phi[i] = 0.5*(phi[i-1]+phi[i+1]+rho_f[i]/p.eps*p.dh*p.dh)
        
#plot results
pl.ion()
fig, ax1 = pl.subplots()

#compute position
r=[]
for p in pieces:
    if len(r)>0:
        r.append(r[-1])
    else: 
        r.append(0)        
    for i in range(1,p.nodes):
        r.append(r[-1]+p.dh)
        
#set permittivity
phi_disp = []
eps_disp = []
for p in pieces:
    for i in range(0,p.nodes):
        eps_disp.append(p.eps)
        phi_disp.append(phi[p.offset+i])
              
ax2 = ax1.twinx()
ax2.plot(r, phi_disp, 'r-')
ax1.plot(r, eps_disp, 'b.', marker='o',markersize=6)

ax1.set_xlabel('position (m)')
ax1.set_ylabel('Eps', color='b')
ax2.set_ylabel('Voltage', color='r')

pl.ioff()

#save notebook
#import IPython.nbformat.current as nbf
#nb = nbf.read(open('dielectric.py','r'),'py')
#nbf.write(nb,open('dielectric.ipynb','w'),'ipynb')