"""Ansatz for variational methods"""

from .spin import generate_In

import numpy as np
import sympy as sp

class Ansatz(object):
    """
    Ansatz is defined in terms of
    
    U_A(theta) = exp[-i(theta/2)A]
    
    as

    |theta_L, ..., theta_1>
    = U_XaYc(theta_L) ... U_YaXc(theta_1) |0, ..., 0>,

    where

    Xa = Xa1 + ... + XaN,
    Ya = Ya1 + ... + YaN

    are collective Pauli operators of the N atomic qubits, 
    and Xc with Yc are Pauli operators for the M=1 cavity qubit.

    One can interprete the role 
    of the atomic qubits and the cavity qubit in the opposite way, 
    i.e., for the case of N=1 atomic qubit and M cavity qubits,
    which corresponds to more natural modeling of the Dicke model,
    since multiple qubits are used to represent the cavity mode.
    """
    def __init__(self, N, zeros_ket=None):
        """
        zeros_ket: (N+1,)
            State vector of |0, ..., 0>
            i.e., ansatz with all angles zero.
        """
        self.M = 1
        self.N = N
        dim = N + 1
        self.dim = dim
        if zeros_ket is None:
            zeros_ket = np.zeros((dim,))
            zeros_ket[0] = 1.
        self.zeros_ket = zeros_ket

        YaXc = sp.zeros(dim)
        for n in range(1,dim):
            YaXc[n,n-1] = sp.I * sp.sqrt(n*(N-n+1))
            YaXc[n-1,n] = sp.conjugate(YaXc[n,n-1])
        self.YaXc = YaXc

        In_list, In_sym_list = generate_In(YaXc, N)

        pn = sp.symbols([fr"p_{{{n}}}" for n in range(dim)], real=True)
        
        I = sp.eye(dim)
        
        U_YX = sp.MutableDenseMatrix(sp.zeros(dim)) 
        for n in range(dim):
            U_YX += pn[n] * pow(-sp.I,n) * In_list[n]

        theta = sp.symbols(r"theta", real=True)
        self.U_YX_cs_func = sp.lambdify([theta] + pn, U_YX, modules='numpy')

    def U_q_func(self, theta, q_is_YX=True):
        N = self.N
        
        c = np.cos(theta/2)
        s = np.sin(theta/2)
        c_pow_n = [1]
        s_pow_n = [1]
        dim = N + 1
        for n in range(1,dim):
            c_pow_n.append(c_pow_n[n-1] * c)
            s_pow_n.append(s_pow_n[n-1] * s)
        args = [c_pow_n[N-n] * s_pow_n[n] for n in range(dim)]
        U_ = self.U_YX_cs_func(theta, *args)
        if not q_is_YX:
            for i in range(2,dim,4):
                U_[i:i+2,:] *= -1
                U_[:,i:i+2] *= -1
        return U_
    
    def __call__(self, *thetas):
        q_is_YX_list = [True, False]
        ket = np.empty((self.dim,), dtype=np.float64)
        ket[:] = self.zeros_ket
        for j, theta_j in enumerate(thetas):
            U = self.U_q_func(theta_j, q_is_YX=q_is_YX_list[j%2])
            ket[:] = U @ ket
        return ket

    def avg_of(self, O):
        thetas_ket = self.ansatz(thetas)
        return np.conj(thetas_ket) @ (O @ thetas_ket)
