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_- |
          ------- 
    
    :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']
        syst[a_lat_neg(sp,0)] = system_params['mu_a']
        syst[b_lat_pos(sp,0)] = system_params['mu_b']
        syst[b_lat_neg(sp,0)] = system_params['mu_b']
        syst[c_lat_pos(sp,0)] = system_params['mu_c']
        syst[c_lat_neg(sp,0)] = system_params['mu_c']
        
        if i < N_c - 1:
        
            
            syst[c_lat_pos(sp+1,0)] = system_params['mu_c']
            syst[b_lat_pos(sp+1,0)] = system_params['mu_b']
            syst[c_lat_neg(sp+1,0)] = system_params['mu_c']
            syst[b_lat_neg(sp+1,0)] = system_params['mu_b']

            # + <--> + 
            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'] #phase e^i pi
            syst[a_lat_neg(sp,0), c_lat_pos(sp+1,0)] = -system_params['j3'] #phase e^i pi
            
        # + <--> + 
        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']
        # - <--> -
        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']
        # + <--> - 
        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'] #phase e^i pi
        syst[a_lat_neg(sp,0), b_lat_pos(sp,0)] = -system_params['j3'] #phase e^i pi
        
    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