"""Hamiltonian"""

from scipy.linalg import eigh_tridiagonal

import numpy as np
import sympy as sp
from numpy import sqrt

class Dicke(object):
    def __init__(self, N, omega, omega0):
        M = 1
        self.M = M
        self.N = N
        self.one_over_sqrt_MN = 1/sqrt(self.M*self.N)

        dim = N + 1
        self.dim = dim
        n = np.arange(dim)
        Zc_diag = np.empty((dim,))
        Zc_diag[0::2] = 1
        Zc_diag[1::2] = -1
        Za_diag = N - 2*n
        self.H0_diag = (-omega/2) * Zc_diag + (-omega0/2) * Za_diag

        self.XaXc_od = np.sqrt(n[1:]*(N-n[1:]+1))

    def H0_avg(self, ket):
        return np.real(np.conj(ket) * self.H0_diag * ket).sum(axis=-1)
        
    def XaXc_avg(self, ket):
        return 2 * np.sum(np.real(np.conj(ket[:-1]) * self.XaXc_od * ket[1:]), axis=-1)

    def H_avg(self, ket, lamb):
        return self.H0_avg(ket) + (lamb * self.one_over_sqrt_MN) * self.XaXc_avg(ket)

    def eigenstate(self, lamb, i=0):
        diag = self.H0_diag
        offidag = (lamb * self.one_over_sqrt_MN) * self.XaXc_od
        eigvals, eigvects = eigh_tridiagonal(diag, offidag, select='i', select_range=(i,i))
        return eigvals[0], eigvects[:,0]

    def __repr__(self):
        return "Spin-Dicke model: H = -omega0 * Za / 2 - omega * Zc / 2 + lambda * XaXc / sqrt(MN)"

