import os
from multiprocessing.connection import Connection
from typing import List

import jax
import pennylane as qml
import pennylane.numpy as np
from jax import numpy as jnp
from pennylane import qchem
from pennylane.devices import Device
from pennylane.optimize import GradientDescentOptimizer

jax.config.update("jax_enable_x64", True)

os.environ["OMP_NUM_THREADS"] = "16"


def parse_xyz_from_text(text: str):
    """Parse XYZ format from text content."""
    lines = text.strip().split("\n")

    # First line: number of atoms
    num_atoms = int(lines[0].strip())

    # Second line: Charge/Multiplicity/Electrons/Orbitals (optional)
    # Skip or parse as needed

    symbols = []
    coordinates = []

    # Parse atom lines (after the second line)
    for line in lines[2 : 2 + num_atoms]:
        parts = line.strip().split()
        if len(parts) >= 4:
            symbol = parts[0]
            x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
            symbols.append(symbol)
            coordinates.append([x, y, z])

    return symbols, coordinates


def extract_electron_info(text: str):
    """Extract electron and orbital counts from the second line."""
    lines = text.strip().split("\n")
    if len(lines) >= 2:
        second_line = lines[1]
        # Parse "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
        electrons = 3  # default
        orbitals = 3  # default
        charge = 0  # default
        multiplicity = 1  # default

        for part in second_line.split():
            if "Electrons=" in part:
                electrons = int(part.split("=")[1])
            elif "Orbitals=" in part:
                orbitals = int(part.split("=")[1])
            elif "Charge=" in part:
                charge = int(part.split("=")[1])
            elif "Multiplicity=" in part:
                multiplicity = int(part.split("=")[1])

        return electrons, orbitals, charge, multiplicity

    return 3, 3, 0, 1  # fallback defaults


def prepare_data(data):
    text_content = data.get("text", "")

    # Parse the molecular data (assuming it's in XYZ format)
    symbols, coordinates = parse_xyz_from_text(text_content)

    # Extract electron/orbital info (from the Charge/Multiplicity line)
    # "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
    electrons, orbitals, charge, multiplicity = extract_electron_info(text_content)

    return {
        "symbols": symbols,
        "coordinates": coordinates,
        "charge": charge,
        "multiplicity": multiplicity,
        "active_electrons": electrons,
        "active_orbitals": orbitals,
        "max_iterations": data.get("max_iterations", 200),
        "conv_tol": data.get("conv_tol", 1e-6),
        "step_size": data.get("step_size", 0.05),
    }


def run_computation(conn: Connection, dev1: Device, data: dict):
    coordinates = jnp.array(data.get("coordinates"))
    charge = int(data.get("charge"))
    multiplicity = int(data.get("multiplicity"))

    molecule = qchem.Molecule(
        data.get("symbols"),
        coordinates,
        charge=charge,
        mult=multiplicity,
    )

    active_electrons = int(data.get("active_electrons"))
    active_orbitals = int(data.get("active_orbitals"))

    max_iterations = int(data.get("max_iterations", 200))
    step_size = float(data.get("step_size", 0.05))
    conv_tol = float(data.get("conv_tol", 1e-6))

    H, qubits = qchem.molecular_hamiltonian(
        molecule,
        active_electrons=active_electrons,
        active_orbitals=active_orbitals,
        method="openfermion",
    )  # type: ignore

    singles, doubles = qml.qchem.excitations(active_electrons, qubits)

    if False:
        params = np.array(last_state, requires_grad=True)
    else:
        params = np.array(np.zeros(len(singles) + len(doubles)), requires_grad=True)

    conn.send(
        {
            "iter_num": 0,
            "energy": None,
            "conv": None,
            "params": params.tolist() if hasattr(params, "tolist") else list(params),
        }
    )

    @qml.qnode(dev1)
    def circuit(param, wires):
        # Map excitations to the wires the UCCSD circuit will act on
        s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)
        qml.UCCSD(
            param,
            wires,
            s_wires=s_wires,
            d_wires=d_wires,
            init_state=qml.qchem.hf_state(active_electrons, qubits),
        )
        return qml.expval(H)

    def cost_fn(param):
        return circuit(param, wires=range(qubits))

    opt = GradientDescentOptimizer(stepsize=step_size)

    for n in range(max_iterations):
        # Take step
        params, prev_energy = opt.step_and_cost(cost_fn, params)

        energy = cost_fn(params)

        # Calculate difference between new and old energies
        conv = np.abs(energy - prev_energy)

        conn.send(
            {
                "iter_num": n,
                "energy": float(energy),
                "conv": float(conv),
                "params": params.tolist()
                if hasattr(params, "tolist")
                else list(params),
            }
        )

        if conv <= conv_tol:
            break