rqism 0.4.1

pub mod qr;

use std::collections::HashMap;

use crate::{
    circuit::Circuit,
    counts::Counts,
    instruction::{CliffordKind, Instruction},
    simulator_traits::Simulator,
};
use ndarray::*;
//use ndarray_linalg::*;
use num::complex::Complex;
use rand::{rngs::ThreadRng, thread_rng, Rng};

pub fn tensor_dot(
    a: &ArrayD<Complex<f32>>,
    b: &ArrayD<Complex<f32>>,
    axes: (Vec<usize>, Vec<usize>),
) -> Result<ArrayD<Complex<f32>>, String> {
    let (axes_a, axes_b) = axes;

    let na = axes_a.len();
    let nb = axes_b.len();

    let as_ = a.shape();
    let nda = a.ndim();
    let bs_ = b.shape();
    let ndb = b.ndim();

    let mut equal = true;

    if na != nb {
        equal = false;
    } else {
        for k in 0..na {
            if as_[axes_a[k]] != bs_[axes_b[k]] {
                equal = false;
                break;
            }
        }
    }

    if !equal {
        return Err("shape mismatch".to_string());
    }

    let notin: Vec<usize> = (0..nda).filter(|k| !axes_a.contains(k)).collect();
    let newaxes_a: Vec<usize> = notin.iter().chain(axes_a.iter()).cloned().collect();
    let n2: usize = axes_a.iter().map(|x| as_[*x]).product();
    let newshape_a = (notin.iter().map(|ax| as_[*ax]).product::<usize>(), n2);
    let olda = notin.iter().map(|a| as_[*a]).collect::<Vec<usize>>();

    let notin: Vec<usize> = (0..ndb).filter(|k| !axes_b.contains(k)).collect();
    let newaxes_b: Vec<usize> = notin.iter().chain(axes_b.iter()).cloned().collect();
    let n2: usize = axes_b.iter().map(|x| bs_[*x]).product();
    let newshape_b = (n2, notin.iter().map(|ax| bs_[*ax]).product::<usize>());
    let oldb = notin.iter().map(|a| bs_[*a]).collect::<Vec<usize>>();

    let at__ = a.view().permuted_axes(newaxes_a);
    //.into_shared()
    let at = at__.to_shape(newshape_a).unwrap();

    let bt__ = b.view().permuted_axes(newaxes_b);
    //.into_shared()
    let bt = bt__.to_shape(newshape_b).unwrap();

    //println!("tensor dot {:?} {:?}", at.shape(), bt.shape());

    Ok(at
        .dot(&bt)
        .into_shared()
        .to_shape(
            olda.iter()
                .chain(oldb.iter())
                .cloned()
                .collect::<Vec<usize>>(),
        )
        .unwrap()
        .into_owned())
}

fn lq_decomposition(a: &Array2<Complex<f32>>) -> (Array2<Complex<f32>>, Array2<Complex<f32>>) {
    let (q, r) = qr::QRDecomp::new(&a.t().to_owned()).unwrap().to_tuple();

    let l = r.t().to_owned(); // L = R^T
    let q_lq = q.t().to_owned(); // Q = Q^T

    (l, q_lq)
}

#[derive(Clone)]
pub struct MPS {
    pub n: usize,
    pub bond_dim: usize,
    pub tensors: Vec<ArrayD<Complex<f32>>>,
    pub rng: ThreadRng,
    pub reg: usize,
}

impl MPS {
    pub fn new(n: usize, bond_dim: usize) -> Self {
        if n == 0 {
            panic!("can not create 0 qubit tensor");
        }

        let tensors: Vec<ArrayD<Complex<f32>>> = (0..n)
            .map(|i| {
                let shape = if i == 0 {
                    (1, 2, bond_dim)
                } else if i == n - 1 {
                    (bond_dim, 2, 1)
                } else {
                    (bond_dim, 2, bond_dim)
                };

                let mut tensor = Array::zeros(shape).into_dyn();
                tensor.slice_mut(s![.., 0, ..]).fill(Complex::new(1.0, 0.0));
                tensor
            })
            .collect();

        let rng = thread_rng();

        Self {
            n,
            bond_dim,
            tensors,
            rng,
            reg: 0,
        }
    }

    pub fn orthonormalize_left(&mut self, i: usize) {
        let shape = self.tensors[i].shape();
        let new_shape = Ix2(shape[0] * shape[1], shape[2]);

        let (q, r) = qr::QRDecomp::new(&self.tensors[i].clone().to_shape(new_shape).unwrap())
            .unwrap()
            .to_tuple();

        // new bond dimension from r
        self.tensors[i] = q
            .to_shape(IxDyn(&[shape[0], shape[1], r.shape()[0]]))
            .unwrap()
            .to_owned();

        let right_tensor_shape = self.tensors[i + 1].shape();

        let absorbed = self.tensors[i + 1]
            .clone()
            .into_shape_with_order([
                right_tensor_shape[0],
                right_tensor_shape[1] * right_tensor_shape[2],
            ])
            .unwrap()
            .dot(&r);

        // new bond dimension from r
        self.tensors[i + 1] = absorbed
            .into_shape_with_order(IxDyn(&[
                r.shape()[0],
                right_tensor_shape[1],
                right_tensor_shape[2],
            ]))
            .unwrap();
    }

    pub fn orthonormalize_right(&mut self, i: usize) {
        let shape = self.tensors[i].shape();
        let new_shape = Ix2(shape[0], shape[1] * shape[2]);

        //println!("{:?}\n{:?}", shape, new_shape);

        let (l, q) = lq_decomposition(
            &self.tensors[i]
                .clone()
                .to_shape(new_shape)
                .unwrap()
                .to_owned(),
        );

        self.tensors[i] = q
            .to_shape(IxDyn(&[l.shape()[0], shape[1], shape[2]]))
            .unwrap()
            .to_owned();

        let left_tensor_shape = self.tensors[i - 1].shape();

        let absorbed = self.tensors[i - 1]
            .clone()
            .into_shape_with_order([
                left_tensor_shape[0] * left_tensor_shape[1],
                left_tensor_shape[2],
            ])
            .unwrap()
            .dot(&l);

        self.tensors[i - 1] = absorbed
            .into_shape_with_order(IxDyn(&[
                left_tensor_shape[0],
                left_tensor_shape[1],
                l.shape()[0],
            ]))
            .unwrap();
    }

    pub fn orthonormalize_around(&mut self, target: usize) {
        // left canoicalization
        for i in 0..target {
            self.orthonormalize_left(i);
        }

        // right canonicalization
        for i in (target + 1..self.tensors.len()).rev() {
            self.orthonormalize_right(i);
        }
    }

    pub fn gate_apply_sq(&mut self, gate: &Array2<Complex<f32>>, i: usize) {
        let contracted = tensor_dot(
            &self.tensors[i],
            &gate.clone().into_dyn(),
            (vec![1], vec![0]),
        )
        .unwrap();

        self.tensors[i] = contracted.permuted_axes(IxDyn(&[0, 2, 1]));
    }

    pub fn expectation_value_sq(
        &self,
        observable: &Array2<Complex<f32>>,
        target: usize,
    ) -> Complex<f32> {
        //expectation value is <psi|G|psi>
        //we do not want to clone all of psi, but in the case of a single qubit,
        //we can only clone one tensor as a time instead of the whole state,
        //then only modify one of the tensors (the target tensor)

        let od = observable.clone().into_dyn();

        let mut acc = Array2::eye(1);

        for i in 0..self.tensors.len() {
            let b = if target == i {
                let contracted = tensor_dot(&self.tensors[i], &od, (vec![1], vec![0])).unwrap();

                contracted.permuted_axes(IxDyn(&[0, 2, 1]))
            } else {
                self.tensors[i].clone()
            };

            let dagger = b.mapv(|x| x.conj());

            let contracted = tensor_dot(&dagger, &self.tensors[i], (vec![1], vec![1])).unwrap();

            let c2d = contracted
                .to_shape((
                    acc.shape()[0] * b.shape()[0] * self.tensors[i].shape()[2],
                    acc.shape()[1] * b.shape()[0] * self.tensors[i].shape()[2],
                ))
                .unwrap();

            acc = acc.dot(&c2d);
            acc = acc.into_shape_with_order((1, 1)).unwrap();
        }

        *acc.get([0_usize, 0]).unwrap()
    }

    pub fn get_reduced_density_matrix(
        &self,
        target: usize,
        normalized: bool,
    ) -> Array2<Complex<f32>> {
        let tensor = if !normalized {
            let mut k = self.clone();
            k.orthonormalize_around(target);

            &k.tensors[target].clone()
        } else {
            &self.tensors[target]
        };

        let shape = tensor.shape();
        let new_shape = (shape[0] * shape[2], shape[1]);

        let m = tensor.to_shape(new_shape).unwrap();

        m.t().mapv(|x| x.conj()).dot(&m)
    }

    pub fn measure_probability(&self, target: usize) -> (f32, f32) {
        let rho = self.get_reduced_density_matrix(target, true);

        //println!("{:?}", rho.shape());

        let p0 = rho.get([0, 0]).unwrap().re;
        let p1 = rho.get([1, 1]).unwrap().re;

        (p0, p1)
    }

    pub fn sample_outcome(&mut self, p0: f32) -> u8 {
        let r = self.rng.gen_range(0.0..1.0_f32);

        if r < p0 {
            0
        } else {
            1
        }
    }

    pub fn project_qubit(&mut self, target: usize, outcome: u8) {
        let phys = 2;
        let bond = 1;

        let mut tensor = Array3::<Complex<f32>>::zeros((bond, phys, bond));
        *tensor.get_mut((0, outcome as usize, 0)).unwrap() = Complex::new(1.0, 0.0);

        self.tensors[target] = tensor.into_dyn();

        // update neighboring bond dimension
        //todo!();
    }

    pub fn measure(&mut self, target: usize, p: f32) -> u8 {
        self.orthonormalize_around(target);
        let (p0, _) = self.measure_probability(target);
        let mut outcome = self.sample_outcome(p0);

        if p != 0.0 && self.rng.gen_range(0.0..1.0) < p {
            outcome = !outcome;
        }

        self.project_qubit(target, outcome);

        if target > 0 {
            self.orthonormalize_left(target - 1);
        }

        if target + 1 < self.tensors.len() {
            self.orthonormalize_right(target + 1);
        }

        self.reg = self.reg & !(1 << target) | (usize::from(outcome) << target);

        outcome
    }

    pub fn measure_state(&mut self) {
        for i in 0..self.tensors.len() {
            self.measure(i, 0.0);
        }
    }

    pub fn execute_instruction(&mut self, ins: &Instruction) {
        match ins {
            Instruction::Pauli { kind, target } => {
                self.gate_apply_sq(&kind.to_matrix(), *target);
            }

            Instruction::Clifford { kind, target } => match kind {
                CliffordKind::H => {
                    self.gate_apply_sq(kind.to_matrix(), target.to_single());
                }
                CliffordKind::S => {
                    self.gate_apply_sq(kind.to_matrix(), target.to_single());
                }
                CliffordKind::CNOT => {
                    panic!("Multi qubit gates not yet implemented");
                    //self.gate_apply_nq(kind.to_matrix(), &target.to_vec());
                }
            },

            Instruction::Identity => {}

            Instruction::Gate {
                matrix, indices, ..
            } => {
                if indices.len() == 1 {
                    self.gate_apply_sq(matrix, indices[0]);
                } else {
                    //self.gate_apply_2q(&matrix, &[indices[0], indices[1]]);
                    //todo!();
                    panic!("Multi qubit gates not yet implemented");
                }
            }

            #[allow(unused_variables)]
            Instruction::Measure { indices } => indices.iter().for_each(|x| {
                self.measure(*x, 0.0);
            }),

            Instruction::MeasureNoisy { indices, p } => indices.iter().for_each(|x| {
                self.measure(*x, *p);
            }),

            Instruction::MeasureState => self.measure_state(),

            Instruction::StateVecMat { .. } => {
                panic!("StateVecMat instructions are only supported on the statevector simulator");
            }

            Instruction::ConditionalInstruction {
                indices,
                condition,
                inst,
            } => {
                let input = indices
                    .iter()
                    .map(|x| ((self.reg >> *x) & 1) != 0)
                    .collect();

                if condition(input) {
                    for ins in inst {
                        self.execute_instruction(&ins);
                    }
                }
            }
        }
    }
    pub fn execute_circuit(&mut self, circuit: &Circuit, index: usize) {
        if index == circuit.ins.len() {
            return;
        }

        self.execute_instruction(&circuit.ins[index]);
        self.execute_circuit(circuit, index + 1);
    }

    pub fn get_counts(&self, circuit: &Circuit, n: usize) -> Counts {
        let mut res = HashMap::new();
        let nq = self.n;

        for _ in 0..n {
            let mut s = self.clone();
            s.execute_circuit(circuit, 0);
            let label = format!("{:0nq$b}", s.reg);

            if let Some(i) = res.get_mut(&label) {
                *i += 1
            } else {
                res.insert(label, 0);
            }
        }

        Counts {
            data: res,
            shots: n,
        }
    }
}

impl Simulator for MPS {
    fn execute_instruction(&mut self, ins: &Instruction) {
        self.execute_instruction(ins);
    }

    fn execute_circuit(&mut self, circuit: &Circuit) {
        self.execute_circuit(circuit, 0);
    }

    fn get_counts(&self, circuit: &Circuit, n: usize) -> Counts {
        self.get_counts(circuit, n)
    }
    fn get_random(&mut self) -> &mut ThreadRng {
        &mut self.rng
    }
}