# -*- coding: utf-8 -*-
import kwant
import numpy as np


def diamond_chain_system(N_c, system_params, semi_infinite = False, leads = False):
    '''
    Create a diamond chain of trimer unit cells. Each atom of the trimer has two orbital angular momentum states, + and -
          _______
    ...  | C_i_+ |
         | C_i_- |         ...
        / ------- \ ______ /
    ...           | A_i_+ | ...
             Φ    | A_i_- |
        \ ______ //------- \
    ...  | B_i_+ |         ...
         | B_i_- |
          ------- 
    
    The Φ represents an out-of-plane magnetic field. The phase is added along the // bond in each unit cell 
    
    :param int N_c: number of unit cells to include in the cell
    :param dict system_params: parameters
    :param bool semi_infinite: whether to make a semi-infinite chain or not
    :param bool leads: whether to include leads 
    
    :rtype kwant.system.FiniteSystem:
    '''
    
    # make lattices and sublattices
    lat = kwant.lattice.Polyatomic(prim_vecs = [[1,0],[0,1]], basis = [[1,0],[1,0], [0,-1],[0,-1], [0,1],[0,1]], norbs = 1)
    a_lat_pos, a_lat_neg, b_lat_pos, b_lat_neg, c_lat_pos, c_lat_neg = lat.sublattices
    
    #make builder and populate with onsite and hoppings
    if semi_infinite == True:
        syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,0]))
    else:
        syst = kwant.Builder()
    
    
    for i in range(N_c):
        
        #staggered point
        sp = i
        
        syst[a_lat_pos(sp,0)] = system_params['mu_a_pos']
        syst[a_lat_neg(sp,0)] = system_params['mu_a_neg']
        syst[b_lat_pos(sp,0)] = system_params['mu_b_pos']
        syst[b_lat_neg(sp,0)] = system_params['mu_b_neg']
        syst[c_lat_pos(sp,0)] = system_params['mu_c_pos']
        syst[c_lat_neg(sp,0)] = system_params['mu_c_neg']
        
        if i < N_c - 1:
        
            
            syst[c_lat_pos(sp+1,0)] = system_params['mu_c_pos']
            syst[b_lat_pos(sp+1,0)] = system_params['mu_b_pos']
            syst[c_lat_neg(sp+1,0)] = system_params['mu_c_neg']
            syst[b_lat_neg(sp+1,0)] = system_params['mu_b_neg']

            # + <--> + 
            syst[a_lat_pos(sp,0), c_lat_pos(sp+1,0)] = system_params['j2']
            syst[a_lat_pos(sp,0), b_lat_pos(sp+1,0)] = system_params['j2']

            # - <--> -
            syst[a_lat_neg(sp,0), c_lat_neg(sp+1,0)] = system_params['j2']
            syst[a_lat_neg(sp,0), b_lat_neg(sp+1,0)] = system_params['j2']

            # + <--> - 
            syst[a_lat_pos(sp,0), b_lat_neg(sp+1,0)] = system_params['j3']
            syst[a_lat_neg(sp,0), b_lat_pos(sp+1,0)] = system_params['j3']

            # + <--> - hopping with phase
            syst[a_lat_pos(sp,0), c_lat_neg(sp+1,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
            syst[a_lat_neg(sp,0), c_lat_pos(sp+1,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
            
        # + <--> + 
        syst[a_lat_pos(sp,0), c_lat_pos(sp,0)] = system_params['j2']
        if i == 0:
            added_phase = 1
        else:
            added_phase = np.exp(1j*2*system_params['phi_d'])
        syst[a_lat_pos(sp,0), b_lat_pos(sp,0)] = system_params['j2']*added_phase #phase e^i phi_d due to mag field
        # - <--> -
        syst[a_lat_neg(sp,0), c_lat_neg(sp,0)] = system_params['j2']
        syst[a_lat_neg(sp,0), b_lat_neg(sp,0)] = system_params['j2']*added_phase #phase e^i phi_d due to mag field
        # + <--> - 
        syst[a_lat_pos(sp,0), c_lat_neg(sp,0)] = system_params['j3']
        syst[a_lat_neg(sp,0), c_lat_pos(sp,0)] = system_params['j3']
        # + <--> - hopping with phase
        syst[a_lat_pos(sp,0), b_lat_neg(sp,0)] = system_params['j3']*np.exp(1j*2*system_params['phi'])*added_phase #phase e^i phi_d due to mag field #phase e^i phi
        syst[a_lat_neg(sp,0), b_lat_pos(sp,0)] = system_params['j3']*np.exp(1j*2*system_params['phi'])*added_phase #phase e^i phi_d due to mag field #phase e^i phi
        
    if leads:
        
        lead_syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,0]))
        
        
        lead_syst[a_lat_pos(0,0)] = 0
        lead_syst[a_lat_neg(0,0)] = 0
        lead_syst[b_lat_pos(0,0)] = 0
        lead_syst[b_lat_neg(0,0)] = 0
        lead_syst[c_lat_pos(0,0)] = 0
        lead_syst[c_lat_neg(0,0)] = 0
        
        lead_syst[lat.neighbors(n=1)] = 1
        
        syst.attach_lead(lead_syst)
        syst.attach_lead(lead_syst.reversed())
    
    return syst

def tilted_diamond_chain_system(l, N_c, system_params, semi_infinite = False, leads = False, closed_chain = False):
    '''
    Create a diamond chain of trimer unit cells. Each atom of the trimer has two orbital angular momentum states, + and -
          
                      _______
              ...  __| B_i_+ |
                     | B_i_- |
                      ------- 
              |     Φ    ||
           _______    _______
          | C_i_+ |__| A_i_+ |__ ...
          | C_i_- |  | A_i_- |
           -------    -------
                          |
                         ...
                         
    The Φ represents an out-of-plane magnetic field. The phase is added along the == bond in each unit cell 
    
    :param int l: the orbital angular momentum number. p bands have l=1, and d have l=2.
    :param int N_c: number of unit cells to include in the cell
    :param dict system_params: parameters
    :param bool semi_infinite: whether to make a semi-infinite chain or not
    :param bool leads: whether to include leads 
    :param bool closed_chain: whether to close the chain by adding one additional A site on the open trimer at one end of the chain.
    
    :rtype kwant.system.FiniteSystem:
    '''
    if l == 0:
        raise ValueError('l cannot be 0: use s_tilted_diamond_chain_system instead.')
    
    # make lattices and sublattices
    lat = kwant.lattice.Polyatomic(prim_vecs = [[1,-1],[1,1]], basis = [[0,0],[0,0], [0,1],[0,1], [-1,0],[-1,0]], norbs = 1)
    a_lat_pos, a_lat_neg, b_lat_pos, b_lat_neg, c_lat_pos, c_lat_neg = lat.sublattices
    
    #make builder and populate with onsite and hoppings
    if semi_infinite == True:
        syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,-1]))
    else:
        syst = kwant.Builder()
    
        
    
    for i in range(N_c):
        
        #staggered point
        sp = i
        
        syst[a_lat_pos(sp,0)] = system_params['mu_a_pos']
        syst[a_lat_neg(sp,0)] = system_params['mu_a_neg']
        syst[b_lat_pos(sp,0)] = system_params['mu_b_pos']
        syst[b_lat_neg(sp,0)] = system_params['mu_b_neg']
        syst[c_lat_pos(sp,0)] = system_params['mu_c_pos']
        syst[c_lat_neg(sp,0)] = system_params['mu_c_neg']
        
        if i == 0:
            
            added_phase = 1
            
            if closed_chain:
                syst[a_lat_pos(-1,0)] = system_params['mu_a_pos']
                syst[a_lat_neg(-1,0)] = system_params['mu_a_neg']
                
                syst[a_lat_pos(-1,0), c_lat_pos(0,0)] = system_params['j2']
                syst[a_lat_pos(-1,0), b_lat_pos(0,0)] = system_params['j2']
                
                syst[a_lat_neg(-1,0), c_lat_neg(0,0)] = system_params['j2']
                syst[a_lat_neg(-1,0), b_lat_neg(0,0)] = system_params['j2']
                
                syst[a_lat_pos(-1,0), b_lat_neg(0,0)] = system_params['j3']*added_phase #phase e^i phi_d due to mag field
                syst[a_lat_neg(-1,0), b_lat_pos(0,0)] = system_params['j3']*added_phase #phase e^i phi_d due to mag field
                
                syst[a_lat_pos(-1,0), c_lat_neg(0,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
                syst[a_lat_neg(-1,0), c_lat_pos(0,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
                
            
        else:
            added_phase = np.exp(1j*2*system_params['phi_d']*l)
            
        
        if i < N_c - 1:
        
            
            syst[c_lat_pos(sp+1,0)] = system_params['mu_c_pos']
            syst[b_lat_pos(sp+1,0)] = system_params['mu_b_pos']
            syst[c_lat_neg(sp+1,0)] = system_params['mu_c_neg']
            syst[b_lat_neg(sp+1,0)] = system_params['mu_b_neg']

            # + <--> + 
            syst[a_lat_pos(sp,0), c_lat_pos(sp+1,0)] = system_params['j2']
            syst[a_lat_pos(sp,0), b_lat_pos(sp+1,0)] = system_params['j2']

            # - <--> -
            syst[a_lat_neg(sp,0), c_lat_neg(sp+1,0)] = system_params['j2']
            syst[a_lat_neg(sp,0), b_lat_neg(sp+1,0)] = system_params['j2']

            # + <--> - 
            syst[a_lat_pos(sp,0), b_lat_neg(sp+1,0)] = system_params['j3']*added_phase #phase e^i phi_d due to mag field
            syst[a_lat_neg(sp,0), b_lat_pos(sp+1,0)] = system_params['j3']*added_phase #phase e^i phi_d due to mag field

            # + <--> - hopping with phase
            syst[a_lat_pos(sp,0), c_lat_neg(sp+1,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
            syst[a_lat_neg(sp,0), c_lat_pos(sp+1,0)] = system_params['j3']*np.exp(1j*2*system_params['phi']) #phase e^i phi
            
        # + <--> + 
        syst[a_lat_pos(sp,0), c_lat_pos(sp,0)] = system_params['j2']
        syst[a_lat_pos(sp,0), b_lat_pos(sp,0)] = system_params['j2']*added_phase #phase e^i phi_d due to mag field
        # - <--> -
        syst[a_lat_neg(sp,0), c_lat_neg(sp,0)] = system_params['j2']
        syst[a_lat_neg(sp,0), b_lat_neg(sp,0)] = system_params['j2']*added_phase #phase e^i phi_d due to mag field
        # + <--> - 
        syst[a_lat_pos(sp,0), c_lat_neg(sp,0)] = system_params['j3']
        syst[a_lat_neg(sp,0), c_lat_pos(sp,0)] = system_params['j3']
        # + <--> - hopping with phase
        syst[a_lat_pos(sp,0), b_lat_neg(sp,0)] = system_params['j3']*np.exp(1j*2*system_params['phi'])*added_phase #phase e^i phi_d due to mag field #phase e^i phi
        syst[a_lat_neg(sp,0), b_lat_pos(sp,0)] = system_params['j3']*np.exp(1j*2*system_params['phi'])*added_phase #phase e^i phi_d due to mag field #phase e^i phi
        
            
        
    if leads:
        
        lead_syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,0]))
        
        
        lead_syst[a_lat_pos(0,0)] = 0
        lead_syst[a_lat_neg(0,0)] = 0
        lead_syst[b_lat_pos(0,0)] = 0
        lead_syst[b_lat_neg(0,0)] = 0
        lead_syst[c_lat_pos(0,0)] = 0
        lead_syst[c_lat_neg(0,0)] = 0
        
        lead_syst[lat.neighbors(n=1)] = 1
        
        syst.attach_lead(lead_syst)
        syst.attach_lead(lead_syst.reversed())
    
    return syst




def s_tilted_diamond_chain_system(N_c, system_params, semi_infinite = False, leads = False, closed_chain = False):
    '''
    Create a diamond chain of trimer unit cells. 
          
                      _______
              ...  __| B_i   |
                     |       |
                      ------- 
              |     Φ    ||
           _______    _______
          | C_i   |__| A_i   |__ ...
          |       |  |       |
           -------    -------
                          |
                         ...
                         
    The Φ represents an out-of-plane magnetic field. The phase is added along the == bond in each unit cell 
    
    :param int N_c: number of unit cells to include in the cell
    :param dict system_params: parameters
    :param bool semi_infinite: whether to make a semi-infinite chain or not
    :param bool leads: whether to include leads 
    :param bool closed_chain: whether to close the chain by adding one additional A site on the open trimer at one end of the chain.
    
    :rtype kwant.system.FiniteSystem:
    '''
    
    # make lattices and sublattices
    lat = kwant.lattice.Polyatomic(prim_vecs = [[1,-1],[1,1]], basis = [[0,0], [0,1], [-1,0]], norbs = 1)
    a_lat, b_lat, c_lat = lat.sublattices
    
    #make builder and populate with onsite and hoppings
    if semi_infinite == True:
        syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,-1]))
    else:
        syst = kwant.Builder()
    
        
    
    for i in range(N_c):
        
        #staggered point
        sp = i
        
        syst[a_lat(sp,0)] = system_params['mu_a']
        syst[b_lat(sp,0)] = system_params['mu_b']
        syst[c_lat(sp,0)] = system_params['mu_c']
        
        if i == 0:
            
            added_phase = 1
            
            if closed_chain:
                syst[a_lat(-1,0)] = system_params['mu_a']
                
                syst[a_lat(-1,0), c_lat(0,0)] = system_params['j2']
                syst[a_lat(-1,0), b_lat(0,0)] = system_params['j2']
                
            
        else:
            added_phase = np.exp(1j*2*system_params['phi_d'])
            
        
        if i < N_c - 1:
        
            #intercell
            syst[c_lat(sp+1,0)] = system_params['mu_c']
            syst[b_lat(sp+1,0)] = system_params['mu_b']

            syst[a_lat(sp,0), c_lat(sp+1,0)] = system_params['j2']
            syst[a_lat(sp,0), b_lat(sp+1,0)] = system_params['j2']

            
        # intracell
        syst[a_lat(sp,0), c_lat(sp,0)] = system_params['j2']
        syst[a_lat(sp,0), b_lat(sp,0)] = system_params['j2']*added_phase #phase e^i phi_d due to mag field
        
        
            
        
    if leads:
        
        lead_syst = kwant.Builder(symmetry=kwant.lattice.TranslationalSymmetry([1,0]))
        
        
        lead_syst[a_lat(0,0)] = 0
        lead_syst[b_lat(0,0)] = 0
        lead_syst[c_lat(0,0)] = 0
        
        lead_syst[lat.neighbors(n=1)] = 1
        
        syst.attach_lead(lead_syst)
        syst.attach_lead(lead_syst.reversed())
    
    return syst