import numpy as np
import matplotlib.pyplot as plt
from .Fockbase.gates import *
from .Fockbase.states import *
from .Fockbase.WignerFunc import *

class Fock():
    def __init__(self, N, cutoff = 10):
        self.N = N
        self.cutoff = cutoff
        self.dim = cutoff + 1
        self.initState = np.zeros([N, self.dim]) + 0j
        self.initState[:, 0] = 1
        self.state = None

    def vacuumState(self, mode):
        #fockState = np.zeros([self.dim ** self.N])
        self.initState[mode, :] = 0
        self.initState[mode, 0] = 1

    def multiTensordot(self):
        self.state = self.initState[0, :]
        for i in range(self.N - 1):
            self.state = np.tensordot(self.state, self.initState[i+1, :], axes = 0)
        return self.state

    def Wigner(self, mode, plot = 'y', xrange = 5.0, prange = 5.0):
        if self.state is None:
            self.state = self.multiTensordot()
            self.initState == None
        x = np.arange(-xrange, xrange, xrange / 50)
        p = np.arange(-prange, prange, prange / 50)
        m = len(x)
        xx, pp = np.meshgrid(x, p)
        xi_array = np.dstack((pp, xx))
        W = FockWigner(xx, pp, self.state, mode)
        if plot == 'y':
            h = plt.contourf(x, p, W)
            plt.show()
        return (x, p, W)
    
    def coherentState(self, mode, alpha):
        if self.initState is None:
            raise ValueError("State must be set before gate operation.")
        self.initState[mode, :] = coherentState(alpha, self.dim)
    
    def catState(self, mode, alpha, parity):
        if self.initState is None:
            raise ValueError("State must be set before gate operation.")
        self.initState[mode, :] = catState(alpha, parity, self.dim)

    def photonNumberState(self, mode, N):
        if self.initState is None:
            raise ValueError("State must be set before gate operation.")
        photonNumState = np.zeros(self.dim)
        photonNumState[N] = 1
        self.initState[mode, :] = photonNumState

    def D(self, mode, alpha):
        if self.state is None:
            self.state = self.multiTensordot()
            self.initState == None
        self.state = displacement(self.state, mode, alpha, self.cutoff)

    def BS(self, mode1, mode2, theta = np.pi/4):
        if self.state is None:
            self.state = self.multiTensordot()
            self.initState == None
        self.state = BS(self.state, mode1, mode2, theta, self.cutoff)

    def S(self, mode, r, phi = 0):
        if self.state is None:
            self.state = self.multiTensordot()
            self.initState == None
        self.state = squeeze(self.state, mode, r, phi, self.cutoff)

    def Kerr(self, mode, chi):
        if self.state is None:
            self.state = self.multiTensordot()
            self.initState == None
        self.state = kerr(self.state, mode, chi, self.cutoff)

    def photonSampling(self, mode, ite = 1):
        reducedDensity = reduceState(self.state, mode)
        probs = np.real(np.diag(reducedDensity))
        probs = probs / np.sum(probs)
        return np.random.choice(probs.shape[0], ite, p = probs)

    def photonMeasurement(self, mode, post_select = None):
        reducedDensity = reduceState(self.state, mode)
        probs = np.real(np.diag(reducedDensity))
        probs = probs / np.sum(probs)
        if post_select is None:
            res = np.random.choice(probs.shape[0], 1, p = probs)
        else:
            res = post_select
        prob = probs[res]
        
        state_ = np.swapaxes(self.state, mode, 0)
        ind = np.ones((self.state.shape[-1]), bool)
        ind[res] = False
        state_[ind] = 0
        state_ = np.swapaxes(state_, mode, 0)
        self.state = state_ / np.sqrt(prob)
        return res