#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 20 13:21:44 2017

@author: alexcurrie1995
"""

import kwant
import semicon
import copy
import combine
import numpy as np
import scipy.linalg as la
import scipy.sparse.linalg as sla

import matplotlib.pyplot as plt

# Define model and system

Gamma_so = [[0, 0, 0, -1],
            [0, 0, +1, 0],
            [0, +1, 0, 0],
            [-1, 0, 0, 0]]

hamiltonian = """
   + M * kron(sigma_0, sigma_z)
   - B * (k_x**2 + k_y**2) * kron(sigma_0, sigma_z)
   - D * (k_x**2 + k_y**2) * kron(sigma_0, sigma_0)
   + A * k_x * kron(sigma_z, sigma_x)
   - A * k_y * kron(sigma_0, sigma_y)
   + Delta * Gamma_so
"""

params_bhz=dict(A=364.5, B=-686, D=-512, M=-10, Delta=1.6, B_x = 0, B_y = 0, B_z = 1.15)

hamiltonian = kwant.continuum.sympify(hamiltonian, locals=dict(Gamma_so=Gamma_so))
hamiltonian

grid_spacing = 5
template = kwant.continuum.discretize(hamiltonian, grid_spacing=grid_spacing)

def apply_peierls_to_template(template, xyz_offset=(0, 0, 0)):
    template = copy.deepcopy(template)  # Needed because kwant.Builder is mutable
    x0, y0, z0 = xyz_offset
    lat = template.lattice
    a = np.max(lat.prim_vecs)  # lattice contant

    def phase(site1, site2, B_x, B_y, B_z, orbital, e, hbar):
        x, y, z = site1.tag
        direction = site2.tag - site1.tag
        A = [B_y * (z - z0) - B_z * (y - y0), 0, B_x * (y - y0)]
        A = np.dot(A, direction) * a**2 * 1e-18 * e / hbar
        phase = np.exp(-1j * A)

        if lat.norbs == 2:  # No PH degrees of freedom
            return phase
        elif lat.norbs == 4:
            return np.array([phase, phase.conj(), phase, phase.conj()],
                            dtype='complex128')

    for (site1, site2), hop in template.hopping_value_pairs():
        template[site1, site2] = combine(hop, phase, operator.mul, 2)

    return template

hbar = 1
#apply_peierls_to_template(template)

syst = kwant.wraparound.wraparound(template)
syst = syst.finalized()

momenta = np.linspace(-np.pi, np.pi, 51)

# Model QPC below.
W = 20
L = 20
H = 15
#H = 250, 270 ,225not as good as 270, 275, 300 not too great.
Gap = W - 2*H
# H = (W - G)/2
def central_shape(site):
    x, y = site.pos
    return 0 < y < W and 0 < x < L and abs(y) < W -  H * np.exp(-(x-L/2)**2 / 1**2) and abs(y) > H * np.exp(-(x-L/2)**2 / 1**2)
    
    

def lead_shape(site):
    x, y = site.pos
    return 0 < y < W

syst = kwant.Builder()
syst.fill(template, central_shape, (grid_spacing, grid_spacing))

lead = kwant.Builder(kwant.TranslationalSymmetry([-grid_spacing, 0]))
lead.fill(template, lead_shape, (0, grid_spacing))

syst.attach_lead(lead)
syst.attach_lead(lead.reversed())

syst = syst.finalized()

energies = np.linspace(-10, 80, 150)
 # Calculate Conductance 
Gs = []
for en in energies:
    smat = kwant.smatrix(syst, energy=en, params=params_bhz)
    Gs.append(smat.transmission(1, 0))
    
y = [0,2,4,6,8,10,12,14,16,18]
plt.plot(energies, Gs)
plt.yticks(np.arange(min(y), max(y)+1, 2.0))
plt.ylim(0, 20)
plt.ylabel("G [$e^2/h$]")
plt.xlabel("energy [meV]")