#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# AUTHOR:     Raj Dahya
# CREATED:    05.09.2022
# DEPARTMENT: Fakult\"at for Mathematik und Informatik
# INSTITUTE:  Universit\"at Leipzig
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# IMPORTS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

import re;
import math
from tkinter import E;
import numpy as np;
from typing import Generator;
import itertools;
from tabulate import tabulate;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# CONSTANTS + SETTINGS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

np.set_printoptions(precision=3);
MACHINE_EPS = 0.5e-12;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# SECONDARY METHODS - Generate C0-Semigroup Generator
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def generate_semigroup_generator(
    shape: list[int, int],
    rational: bool = False,
    base: int = 2,
    alpha: float = 1,
) -> np.ndarray:
    '''
    Creates bounded generators for d commuting C0-semigroups as in the paper (cf. §5.3).

    @inputs
    - `shape`    - <int, int> Desired shape of output.
    - `rational` - <bool>     Whether or not entries are to be rational.
    - `base`     - <int>      If `rational = True`, fixes the denominator of the rational numbers.
    - `alpha`    - <float>    Additional parameter to scale the D_i operators.

    NOTE: in the paper `α = 1` was chosen. However one can choose any value in `(1/√d, \infty)`.
    By choosing any value `α ≥ 1/√(d-1)`, by the computations in Proposition 5.3
    one can force that the S_TK operators only fail to be positive when |K| > d-1.

    @returns
    Randomly created bounded generators for C0-semigroups as in the paper (cf. §5.3).
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    assert m % 2 == 0, 'dimensions must be divisible by 2';
    N = int(m/2);
    I = generate_op_identity(shape=[2*N, 2*N]);
    O = generate_op_zero(shape=[N, N]);
    U = generate_op_unitary(
        shape = [N, N],
        rational = rational,
        base = base,
    );
    D = alpha * U;
    A = -I + np.vstack([
        np.hstack([ O, -2*D ]),
        np.hstack([ O,  O  ]),
    ]);
    return A;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# SECONDARY METHODS - Compute spectral bounds
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def spec_bounds(A: np.ndarray) -> float:
    '''
    Computes the spectral bounds of an operator:
    ```tex
    \max_{i} \max \{ Re \lambda : \lambda \in \sigma(A_{i}) \}
    ```
    '''
    sigma_A = np.linalg.eigvals(A);
    return max(np.real(sigma_A));

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# SECONDARY METHODS - Generate Dissipation Operator
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def dissipation_operators(
    shape: list[int, int],
    A: list[np.ndarray],
) -> tuple[
    list[tuple[
        list[int],
        np.ndarray,
        float,
    ]],
    float,
]:
    '''
    Computes the dissipation operators and their minimum spectral values.
    '''
    d = len(A);
    indexes = list(range(d));
    data = [];
    beta = np.inf;
    for K in power_set(indexes):
        S_TK, eig_K = dissipation_operator(shape=shape, A=A, K=K);
        data.append((K, S_TK, eig_K));
        beta = min(eig_K, beta);
    return data, beta;

def dissipation_operator(
    shape: list[int, int],
    A: list[np.ndarray],
    K: list[int],
) -> tuple[np.ndarray, float]:
    '''
    Computes the dissipation operator `S_{T,K}` and its minimum spectral values.
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    N = m;
    O = generate_op_zero(shape=[N, N]);
    I = generate_op_identity(shape=[N, N]);
    S = O;
    for C, C_ in partitions_two(K):
        A_C = I;
        A_C_ = I;
        for i in C:
            A_C = A_C @ -A[i];
        for i in C_:
            A_C_ = A_C_ @ -A[i];
        S += A_C.conj().T @ A_C_;
    S = (1/2**len(K)) * S;
    eig = min(np.real(np.linalg.eigvals(S)));
    return S, eig;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# MISCELLANEOUS METHODS - SETS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def power_set(K: list[int]) -> Generator[list[int], None, None]:
    '''
    Iterates through all subsets `C` of a set `K`.
    '''
    n = len(K);
    for ch in itertools.product(*([[0,1]]*n)):
        yield [ K[i] for i in range(n) if ch[i] == 1 ];

def partitions_two(K: list[int]) -> Generator[tuple[list[int],list[int]], None, None]:
    '''
    Iterates through all partitions `[C1, C2]` of a set `K` of size 2.
    '''
    n = len(K);
    for ch in itertools.product(*([[0,1]]*n)):
        yield (
            [ K[i] for i in range(n) if ch[i] == 1 ],
            [ K[i] for i in range(n) if ch[i] == 0 ],
        );

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# MISCELLANEOUS METHODS - BASIC OPERATORS
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def generate_op_zero(shape: list[int, int]) -> np.ndarray:
    '''
    Returns the zero operator with the right shape and type.
    '''
    sh = tuple(shape)
    O = np.zeros(sh).astype(complex);
    return O;

def generate_op_identity(shape: list[int, int]) -> np.ndarray:
    '''
    Returns the identity operator with the right shape and type.
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    I = np.eye(m).astype(complex);
    return I;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# MISCELLANEOUS METHODS - GENERATE RANDOM MATRICES
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def generate_op_selfadjoint(
    shape: list[int, int],
    bound: float = 1.,
) -> np.ndarray:
    '''
    Generates a self-adjoint operator.

    @inputs
    - `shape` - <int, int> Desired shape of the output.
    - `bound` - <float>    The desired norm of output operator.

    @returns
    A randomly generated self-adjoint operator as a matrix.
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    A = np.random.random(sh) + 1j * np.random.random(sh);
    S = (A + A.conj().T)/2.;
    M = np.linalg.norm(S);
    if M > 0:
        S = (bound/M) * S;
    return S;

def generate_onb(shape: list[int, int]) -> np.ndarray:
    '''
    Generates an ONB.

    @inputs
    - `shape` - <int, int> Desired shape of matrices.

    @returns
    An ONB for an n-dim complex Hilbert space as a matrix.
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    N = m;
    O = generate_op_zero(shape=[N, N]);
    P = O;
    for k in range(N):
        while True:
            v = np.random.random((N,)) + 1j * np.random.random((N,));
            for i in range(k):
                # NOTE: np.inner does not work.
                coeff = np.sum([ v[j] * P[j, i].conj() for j in range(N)]);
                v -= coeff * P[:, i];
            norm_v = np.linalg.norm(v);
            if norm_v > 0.:
                P[:, k] = v / norm_v;
                break;
    return P;

def generate_op_unitary(
    shape: list[int, int],
    rational: bool = False,
    base: int = 2,
) -> np.ndarray:
    '''
    Generates a unitary operator.

    @inputs
    - `shape`    - <int, int> Desired shape of output.
    - `rational` - <bool>     Whether or not entries are to be rational.
    - `base`     - <int>      If `rational = True`, fixes the denominator of the rational numbers.

    @returns
    A randomly generated unitary operator as a matrix.
    '''
    sh = tuple(shape)
    m, n = sh;
    assert m == n, 'must be square';
    N = m;
    P = generate_onb(shape=shape);
    if rational:
        theta = 2 * np.pi * np.random.randint(low=0, high=base, size=N) / base;
    else:
        theta = 2 * np.pi * np.random.random(size=N);
    D = np.diag(np.exp(1j * theta));
    U = P @ D @ P.conj().T;
    return U;

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# MISCELLANEOUS METHODS - Get user input
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def ask_confirmation(message: str):
    while True:
        answer = input(f'{message} (y/n): ');
        if re.match(pattern='^y(es)?|1$', string=answer, flags=re.IGNORECASE):
            return True;
        elif re.match(pattern='^n(o)?|0$', string=answer, flags=re.IGNORECASE):
            return False;