"""Spin rotation"""

import sympy as sp

def generate_In(A_mat, N):

    dim = A_mat.shape[0]
    
    A, N_sym, I_sym, Aj = sp.symbols(r"A,N,I,A_{j}")
    
    subs_j = [
        (A,A-Aj),
        (N_sym,N_sym-1)
    ]
    
    subs_order_lowering = []
    
    In_list = [sp.Integer(1), A]
    
    for n in range(2,dim):
        subs_order_lowering = [(Aj**n,Aj**(n%2))] + subs_order_lowering
        Aj_I_n_minus_1_j = (Aj * In_list[n-1].subs(subs_j)).expand().subs(subs_order_lowering)
        
        Aj_I_n_minus_1_j_const = Aj_I_n_minus_1_j.subs(Aj,0)
        In_sym = sp.Rational(1,n) * sp.expand(Aj_I_n_minus_1_j.subs(Aj,A) - Aj_I_n_minus_1_j_const + Aj_I_n_minus_1_j_const * N_sym)
    
        In_list.append(In_sym)
    
    # [NOTE] this loop should be performed after the previous loop ends
    for n, In_sym in enumerate(In_list):
        In_sym_zero_A = In_sym.subs(A,0)
        In_sym = In_sym - In_sym_zero_A + In_sym_zero_A * I_sym
        In_list[n] = In_sym
    
    I = sp.eye(dim)
    
    subs_matr = [
        (I_sym, I),
        (A, A_mat),
        (N_sym, N)
    ]

    In_matr_list = []
    for n, In_sym in enumerate(In_list):
        In_matr = In_sym.subs(subs_matr)
        In_matr_list.append(In_matr)
        
    return In_matr_list, In_list

