"""Single-site bosonic and fermionic operator factories."""

import numpy as np
from scipy.sparse import csr_matrix, diags
from scipy.sparse import identity
from edlgt.modeling import qmb_operator as qmb_op

__all__ = [
    "fermi_operators",
    "bose_operators",
]


def fermi_operators(has_spin, colors=False, fermionic=True):
    """Build single-site fermionic matter operators.

    Parameters
    ----------
    has_spin : bool
        If ``True``, build two-component (up/down) matter operators.
    colors : bool, optional
        If ``True``, provide ``r/g`` aliases for the two matter components.
    fermionic : bool, optional
        If ``True``, include fermionic parity operators/sign strings.

    Returns
    -------
    dict
        Dictionary of single-site matter operators.
    """
    ops = {}
    # Define the MATTER FIELDS OPERATORS
    ops["psi"] = diags(np.array([1], dtype=float), 1, (2, 2))
    ops["psi_dag"] = ops["psi"].transpose()
    if fermionic:
        ops["P_psi"] = diags(np.array([1, -1], dtype=float), 0, (2, 2))
    else:
        ops["P_psi"] = identity(2, dtype=float)
    ops["N"] = ops["psi_dag"] * ops["psi"]
    ops["ID_psi"] = identity(2, dtype=float)
    ops["ID"] = identity(2, dtype=float)
    ops["psi_dag_P"] = ops["psi_dag"] @ ops["P_psi"]
    if has_spin:
        # up & down MATTER OPERATORS
        ops["psi_down"] = qmb_op(ops, ["psi", "ID"])
        ops["psi_up"] = qmb_op(ops, ["P_psi", "psi"])
        ops["N_down"] = qmb_op(ops, ["N", "ID"])
        ops["N_up"] = qmb_op(ops, ["ID", "N"])
        # other number operators
        ops["N_pair"] = ops["N_up"] * ops["N_down"]
        ops["N_tot"] = ops["N_up"] + ops["N_down"]
        ops["N_single"] = ops["N_tot"] - 2 * ops["N_pair"]
        ops["N_zero"] = identity(4, dtype=float) - ops["N_pair"] - ops["N_single"]
        # identity on the whole matter site
        ops["ID_psi"] = identity(4, dtype=float)
        ops["P_psi"] = qmb_op(ops, ["P_psi", "P_psi"])
        for s in ["up", "down"]:
            ops[f"psi_{s}_dag"] = ops[f"psi_{s}"].transpose()
            ops[f"psi_{s}_dag_P"] = ops[f"psi_{s}_dag"] @ ops["P_psi"]
        if colors:
            ops["psi_r"] = ops["psi_up"]
            ops["psi_g"] = ops["psi_down"]
            ops["N_r"] = ops["N_up"]
            ops["N_g"] = ops["N_down"]
            for s in ["r", "g"]:
                ops[f"psi_{s}_dag"] = ops[f"psi_{s}"].transpose()
                ops[f"psi_{s}_dag_P"] = ops[f"psi_{s}_dag"] @ ops["P_psi"]
    else:
        ops["N_zero"] = ops["ID"] - ops["N"]
    return ops


def bose_operators(n_max):
    """Build bosonic ladder and number operators with onsite cutoff ``n_max``.

    Parameters
    ----------
    n_max : int
        Maximum onsite occupation number.

    Returns
    -------
    dict
        Dictionary containing ``b``, ``b_dagger``, ``N``, and ``N2``.
    """
    ops = {}
    entries = np.arange(1, n_max + 1, 1)
    entries = np.sqrt(entries)
    x_coords = np.arange(0, n_max, 1)
    y_coords = np.arange(1, n_max + 1, 1)
    ops["b"] = csr_matrix((entries, (x_coords, y_coords)), shape=(n_max + 1, n_max + 1))
    ops["b_dagger"] = csr_matrix(ops["b"].conj().transpose())
    ops["N"] = ops["b_dagger"] * ops["b"]
    ops["N2"] = ops["N"] ** 2
    return ops
