"""
BIOOS Alpha-Flex Fano-Basin Engine
Ported from BIOOS V48: THE SOVEREIGN APEX
Author: Charles Tibedo | 2026

Core physics engine for protein sequence analysis via golden-ratio
Fano-coordinate manifold mapping. Numba-optimized.
"""

import math
import random
import numpy as np
from functools import lru_cache

try:
    from numba import njit
    HAS_NUMBA = True
except ImportError:
    HAS_NUMBA = False

# ═══════════════════════════════════════════════════════════════
# CONSTANTS
# ═══════════════════════════════════════════════════════════════

PHI = (1.0 + 5.0**0.5) / 2.0  # Golden ratio φ = 1.618034...

# Hydrophobicity / physicochemical scales (Kyte-Doolittle extended)
ALPHABET_SCALES = {
    'R': -4.5, 'K': -3.9, 'N': -3.5, 'D': -3.5, 'Q': -3.5, 'E': -3.5,
    'H': -3.2, 'P': -1.6, 'Y': -1.3, 'W': -0.9, 'S': -0.8, 'T': -0.7,
    'G': -0.4, 'A': 1.8, 'M': 1.9, 'C': 2.5, 'F': 2.8, 'L': 3.8,
    'V': 4.2, 'I': 4.5
}

# Fano coordinates for 8-point manifold mapping
# Point 0 = origin, Points 1-7 = unit-normalized Fano vertices
FANO_COORDS = np.array([
    [0.0, 0.0, 0.0],
    [1.0, 1.0, 1.0],
    [-1.0, -1.0, 1.0],
    [-1.0, 1.0, -1.0],
    [1.0, -1.0, -1.0],
    [0.0, 0.0, 1.5],
    [0.0, 1.5, 0.0],
    [1.5, 0.0, 0.0]
], dtype=np.float64)

# Normalize points 1-7 to unit sphere
for i in range(1, 8):
    norm = np.linalg.norm(FANO_COORDS[i])
    if norm > 0:
        FANO_COORDS[i] /= norm

# 23-valve SPPS automation mapping
VALVE_MAP = {
    'A': 'V01', 'C': 'V02', 'D': 'V03', 'E': 'V04', 'F': 'V05',
    'G': 'V06', 'H': 'V07', 'I': 'V08', 'K': 'V09', 'L': 'V10',
    'M': 'V11', 'N': 'V12', 'P': 'V13', 'Q': 'V14', 'R': 'V15',
    'S': 'V16', 'T': 'V17', 'V': 'V18', 'W': 'V19', 'Y': 'V20',
    'DMF': 'V21', 'PIPERIDINE': 'V22', 'ACTIVATOR': 'V23'
}

# Per-residue coupling times (seconds)
COUPLING_TIMES = {
    'P': 3600,   # Proline: slow coupling (steric)
    'W': 2400,   # Aromatics: extended coupling
    'Y': 2400,
    'F': 2400,
}
DEFAULT_COUPLING_TIME = 1800


# ═══════════════════════════════════════════════════════════════
# ALPHA-FLEX STEP (Pure Python fallback + Numba when available)
# ═══════════════════════════════════════════════════════════════

def _alpha_flex_step_py(scalar, biases, h_coh, psr, psi):
    """
    Core Alpha-Flex computation step.

    Parameters:
        scalar: hydrophobicity value for current residue
        biases: 8-element Fano bias array
        h_coh: harmonic coherence factor
        psr: previous signal real component
        psi: previous signal imaginary component

    Returns:
        (sigma, determinant, new_sr, new_si)
    """
    z_mag = PHI ** scalar

    # Fano lane selection via arctangent modular arithmetic
    arg1 = abs(z_mag * math.pi - round(z_mag * math.pi))
    arg2 = abs(z_mag - PHI * round(scalar))
    arg2 = arg2 if arg2 > 1e-18 else 1e-9
    k = int((math.atan2(arg1, arg2) * 14.0 / math.pi)) % 7

    bias = biases[k]

    # Harmonic theta with coherence modulation
    h_theta = (k * math.pi / 3.0) + (h_coh * 1.745)

    # Complex signal propagation
    sr = (z_mag * bias * math.cos(h_theta)) + psr
    si = (z_mag * bias * math.sin(h_theta)) + psi

    # Magnitude and safe division
    mag = math.sqrt(sr * sr + si * si)
    safe = mag if mag > 1e-18 else 1e-9

    # Sigma: golden ratio / pi proximity metric
    phi_dist = abs(safe - PHI * round(math.log(safe + 1e-18) / math.log(PHI)))
    pi_dist = abs(safe * math.pi - round(safe * math.pi))
    sig = phi_dist * pi_dist

    # Determinant: stability metric with exponential damping
    grad = abs((PHI ** (scalar + 1e-5) - PHI ** (scalar - 1e-5)) / 2e-5)
    det = grad * math.exp(-sig) * (1.0 + 0.5 * math.cos(h_coh))

    return sig, det, sr, si


# Use Numba-compiled version if available
if HAS_NUMBA:
    @njit(fastmath=True)
    def alpha_flex_step(scalar, biases, h_coh, psr, psi):
        z_mag = PHI ** scalar
        arg1 = abs(z_mag * math.pi - round(z_mag * math.pi))
        arg2 = abs(z_mag - PHI * round(scalar))
        if arg2 < 1e-18:
            arg2 = 1e-9
        k = int((math.atan2(arg1, arg2) * 14.0 / math.pi)) % 7
        bias = biases[k]
        h_theta = (k * math.pi / 3.0) + (h_coh * 1.745)
        sr = (z_mag * bias * math.cos(h_theta)) + psr
        si = (z_mag * bias * math.sin(h_theta)) + psi
        mag = math.sqrt(sr * sr + si * si)
        safe = mag if mag > 1e-18 else 1e-9
        phi_dist = abs(safe - PHI * round(math.log(safe + 1e-18) / math.log(PHI)))
        pi_dist = abs(safe * math.pi - round(safe * math.pi))
        sig = phi_dist * pi_dist
        grad = abs((PHI ** (scalar + 1e-5) - PHI ** (scalar - 1e-5)) / 2e-5)
        det = grad * math.exp(-sig) * (1.0 + 0.5 * math.cos(h_coh))
        return sig, det, sr, si
else:
    alpha_flex_step = _alpha_flex_step_py


# ═══════════════════════════════════════════════════════════════
# QUANTUM ASSEMBLER
# ═══════════════════════════════════════════════════════════════

class QuantumAssembler:
    """
    Runs Alpha-Flex chains on protein sequences and returns
    manifold trajectories, signal arrays, and determinant metrics.
    """

    def __init__(self):
        self._cache = {}

    def run_chain(self, seq, biases, flex=1.0):
        """
        Execute Alpha-Flex chain on a protein sequence.

        Parameters:
            seq: amino acid sequence string
            biases: numpy array of 8 Fano bias values
            flex: Alpha-Flex coherence multiplier (Omega_H)

        Returns:
            dict with mean_det, mean_sig, sigs, dets, manifold, steps
        """
        ckey = (seq, tuple(np.round(biases, 4)), round(flex, 4))
        if ckey in self._cache:
            return self._cache[ckey]

        n = len(seq)
        if n == 0:
            return {"mean_det": 0, "mean_sig": 0, "sigs": [], "dets": [], "manifold": [], "steps": []}

        sigs = np.zeros(n)
        dets = np.zeros(n)
        steps = np.zeros(n, dtype=np.complex128)

        # Map sequence to hydrophobicity scalars
        scs = np.array([ALPHABET_SCALES.get(a, 0.0) for a in seq])
        sr, si = 0.0, 0.0

        for i in range(n):
            h_coh = math.sin(i * 1.745) * flex
            sigs[i], dets[i], sr, si = alpha_flex_step(scs[i], biases, h_coh, sr, si)
            steps[i] = complex(sr, si)

        m_sig = np.mean(sigs)

        # Cumulative manifold trajectory with signal-weighted damping
        weights = np.where(sigs < m_sig, 1.0, np.exp(-(sigs - m_sig)))
        Z = np.cumsum(steps * weights)

        res = {
            "mean_det": float(np.mean(dets)),
            "mean_sig": float(m_sig),
            "max_det": float(np.max(dets)),
            "min_sig": float(np.min(sigs)),
            "sigs": sigs.tolist(),
            "dets": dets.tolist(),
            "manifold_real": np.real(Z).tolist(),
            "manifold_imag": np.imag(Z).tolist(),
            "fano_lanes": self._compute_fano_lanes(scs),
        }

        # Cache management
        if len(self._cache) > 5000:
            self._cache.clear()
        self._cache[ckey] = res
        return res

    def _compute_fano_lanes(self, scalars):
        """Compute which Fano lane (0-6) each residue maps to."""
        lanes = []
        for s in scalars:
            z_mag = PHI ** s
            arg1 = abs(z_mag * math.pi - round(z_mag * math.pi))
            arg2 = abs(z_mag - PHI * round(s))
            arg2 = arg2 if arg2 > 1e-18 else 1e-9
            k = int((math.atan2(arg1, arg2) * 14.0 / math.pi)) % 7
            lanes.append(k)
        return lanes

    def compute_interactome(self, sequences, biases, flex=1.0):
        """
        Compute pairwise interactome matrix for multiple protein sequences.

        Returns NxN matrix of signal/determinant cross-correlation coefficients.
        """
        chains = [self.run_chain(seq, biases, flex) for seq in sequences]
        n = len(chains)
        matrix = np.zeros((n, n))

        for i in range(n):
            for j in range(n):
                diff_sig = abs(chains[i]['mean_sig'] - chains[j]['mean_sig'])
                matrix[i, j] = (chains[i]['mean_det'] * chains[j]['mean_det']) * math.exp(-diff_sig)

        return {
            "matrix": matrix.tolist(),
            "chains": chains,
        }

    def compute_fano_tensor(self, biases):
        """
        Compute bias-modulated Fano coordinate tensor.
        Returns 8 vectors from origin, scaled by biases.
        """
        padded = np.ones(8)
        padded[:len(biases)] = biases[:8]
        vectors = FANO_COORDS * padded[:, np.newaxis]
        return {
            "vectors": vectors.tolist(),
            "magnitudes": [float(np.linalg.norm(v)) for v in vectors],
            "biases": padded.tolist(),
        }


# ═══════════════════════════════════════════════════════════════
# GENETIC ALGORITHM SEQUENCE OPTIMIZER
# ═══════════════════════════════════════════════════════════════

class SequenceOptimizer:
    """
    Genetic algorithm for de novo peptide sequence evolution.
    Applies aromatic stacking rules and Alpha-Flex fitness scoring.
    """

    def __init__(self, engine):
        self.engine = engine
        self.amino_acids = list(ALPHABET_SCALES.keys())

    def optimize(self, biases, flex=1.0, seq_length=15, pop_size=100,
                 generations=100, top_k=10, mutation_rate=0.1):
        """
        Evolve optimal peptide sequences.

        Returns:
            dict with best sequence, fitness trajectory, top candidates
        """
        biases = np.array(biases)

        # Initialize random population
        population = [
            "".join(random.choices(self.amino_acids, k=seq_length))
            for _ in range(pop_size)
        ]

        fitness_history = []
        all_top = []

        for gen in range(generations):
            scored = []
            for seq in population:
                r = self.engine.run_chain(seq, biases, flex)

                # Aromatic stacking bonus (W/Y/F at i+3 and i+4)
                bonus = 1.0
                aromatics = {'W', 'Y', 'F'}
                for i in range(len(seq) - 4):
                    if seq[i] in aromatics:
                        if seq[i + 3] in aromatics:
                            bonus += 0.25
                        if seq[i + 4] in aromatics:
                            bonus += 0.25

                fitness = r['mean_det'] * math.exp(-r['mean_sig']) * bonus
                scored.append((fitness, seq))

            # Sort by fitness (descending)
            scored.sort(reverse=True, key=lambda x: x[0])

            # Track history
            fitness_history.append({
                "generation": gen,
                "best_fitness": scored[0][0],
                "mean_fitness": sum(s[0] for s in scored) / len(scored),
                "best_seq": scored[0][1],
            })

            # Keep top-k
            elite = [s[1] for s in scored[:top_k]]

            # Reproduce with crossover + mutation
            population = list(elite)
            while len(population) < pop_size:
                p1, p2 = random.sample(elite, 2)
                split = random.randint(1, seq_length - 1)
                child = p1[:split] + p2[split:]

                # Point mutation
                if random.random() < mutation_rate:
                    idx = random.randint(0, seq_length - 1)
                    child = child[:idx] + random.choice(self.amino_acids) + child[idx + 1:]

                population.append(child)

        # Final scoring
        final_scored = []
        for seq in population[:top_k]:
            r = self.engine.run_chain(seq, biases, flex)
            bonus = 1.0
            aromatics = {'W', 'Y', 'F'}
            for i in range(len(seq) - 4):
                if seq[i] in aromatics:
                    if seq[i + 3] in aromatics:
                        bonus += 0.25
                    if seq[i + 4] in aromatics:
                        bonus += 0.25
            fitness = r['mean_det'] * math.exp(-r['mean_sig']) * bonus

            # Count aromatic stacks
            stacks = 0
            for i in range(len(seq) - 4):
                if seq[i] in aromatics:
                    if seq[i + 3] in aromatics:
                        stacks += 1
                    if seq[i + 4] in aromatics:
                        stacks += 1

            final_scored.append({
                "sequence": seq,
                "fitness": round(fitness, 6),
                "mean_det": round(r['mean_det'], 6),
                "mean_sig": round(r['mean_sig'], 6),
                "aromatic_stacks": stacks,
                "fano_lanes": r['fano_lanes'],
            })

        final_scored.sort(key=lambda x: x['fitness'], reverse=True)

        return {
            "best": final_scored[0] if final_scored else None,
            "top_candidates": final_scored,
            "fitness_history": fitness_history,
            "parameters": {
                "seq_length": seq_length,
                "pop_size": pop_size,
                "generations": generations,
                "flex": flex,
                "mutation_rate": mutation_rate,
            }
        }


# ═══════════════════════════════════════════════════════════════
# G-CODE SPPS COMPILER
# ═══════════════════════════════════════════════════════════════

class SPPSCompiler:
    """
    Converts validated peptide sequences into automated SPPS G-Code
    for robotic synthesis hardware. C-to-N reversed synthesis with
    per-residue timing logic.
    """

    @staticmethod
    def compile(sequence, job_id="MANUAL", flex=1.0, include_header=True):
        """
        Generate G-Code for SPPS robotic synthesis.

        Parameters:
            sequence: amino acid sequence (N-to-C)
            job_id: identifier for the synthesis job
            flex: Omega_H coherence parameter

        Returns:
            G-Code string
        """
        # C-to-N reversal for Fmoc SPPS
        reversed_seq = sequence[::-1]

        lines = []
        if include_header:
            lines.extend([
                f"; ═══════════════════════════════════════════",
                f"; BIOOS V48 SPPS G-CODE",
                f"; JOB: {job_id}",
                f"; SEQUENCE: {sequence}",
                f"; LENGTH: {len(sequence)} residues",
                f"; OMEGA_H: {flex}",
                f"; SYNTHESIS: C→N (Fmoc/tBu)",
                f"; ═══════════════════════════════════════════",
                "",
                "; --- SYSTEM INITIALIZATION ---",
                "ACTION:INIT VALVE:SYSTEM",
                "ACTION:DISPENSE VALVE:V21 VOL:5.0mL ; DMF wash",
                "ACTION:DRAIN",
                "",
            ])

        total_time = 0
        for i, aa in enumerate(reversed_seq):
            coupling_time = COUPLING_TIMES.get(aa, DEFAULT_COUPLING_TIME)
            total_time += coupling_time + 180  # coupling + deprotection
            valve = VALVE_MAP.get(aa, 'V01')

            lines.extend([
                f"; --- CYCLE {i+1}/{len(sequence)}: {aa} (Valve {valve}) ---",
                f"ACTION:DISPENSE VALVE:V22 VOL:2.0mL ; Fmoc deprotection (20% piperidine/DMF)",
                f"ACTION:AGITATE TIME:180s",
                f"ACTION:DRAIN",
                f"ACTION:DISPENSE VALVE:V21 VOL:3.0mL ; DMF wash x3",
                f"ACTION:DRAIN",
                f"ACTION:DISPENSE VALVE:{valve} VOL:2.0mL ; Couple {aa}",
                f"ACTION:DISPENSE VALVE:V23 VOL:1.0mL ; HBTU/DIPEA activator",
                f"ACTION:AGITATE TIME:{coupling_time}s",
                f"ACTION:DRAIN",
                f"ACTION:DISPENSE VALVE:V21 VOL:3.0mL ; DMF wash",
                f"ACTION:DRAIN",
                "",
            ])

        lines.extend([
            "; --- CLEAVAGE & DEPROTECTION ---",
            "ACTION:DISPENSE VALVE:SYSTEM VOL:10.0mL ; TFA/TIS/H2O (95:2.5:2.5)",
            "ACTION:AGITATE TIME:7200s ; 2hr cleavage",
            "ACTION:COLLECT",
            "",
            "; --- FINALIZATION ---",
            "ACTION:FINALIZE VALVE:SYSTEM",
            f"; TOTAL ESTIMATED TIME: {total_time}s ({total_time/3600:.1f}hr)",
            f"; RESIDUE COUNT: {len(sequence)}",
            f"; END OF JOB {job_id}",
        ])

        return "\n".join(lines)

    @staticmethod
    def compile_batch(sequences, job_prefix="BATCH"):
        """Compile multiple sequences into a single G-Code file."""
        all_codes = []
        for i, seq in enumerate(sequences):
            code = SPPSCompiler.compile(
                seq,
                job_id=f"{job_prefix}_{i+1:03d}",
                include_header=(i == 0)
            )
            all_codes.append(code)
            all_codes.append("\n; ═══════════════════════════════════════════\n")
        return "\n".join(all_codes)


# ═══════════════════════════════════════════════════════════════
# BIAS OPTIMIZER (Nelder-Mead)
# ═══════════════════════════════════════════════════════════════

class BiasOptimizer:
    """
    Optimizes the 8-element Fano bias vector to maximize
    mean determinant across a set of target sequences.
    """

    def __init__(self, engine):
        self.engine = engine

    def optimize(self, sequences, initial_biases=None, flex=1.0, max_iter=200):
        """
        Run Nelder-Mead optimization on Fano biases.

        Returns:
            dict with optimized biases, fitness trajectory, convergence info
        """
        from scipy.optimize import minimize as sp_minimize

        if initial_biases is None:
            initial_biases = np.ones(8)
        else:
            initial_biases = np.array(initial_biases)

        trajectory = []

        def objective(x):
            results = []
            for seq in sequences:
                # Cap at 200 residues for optimization performance
                r = self.engine.run_chain(seq[:200], x, flex)
                results.append(r['mean_det'])
            fitness = np.mean(results)
            trajectory.append({
                "biases": x.tolist(),
                "fitness": float(fitness),
            })
            return 1.0 / (fitness + 1e-9)

        result = sp_minimize(
            objective,
            initial_biases,
            method='Nelder-Mead',
            options={'maxiter': max_iter, 'xatol': 1e-5, 'fatol': 1e-5}
        )

        return {
            "optimized_biases": result.x.tolist(),
            "final_fitness": float(1.0 / (result.fun + 1e-9)),
            "iterations": result.nit,
            "converged": bool(result.success),
            "message": result.message,
            "trajectory": trajectory[-20:],  # Last 20 steps
        }


# ═══════════════════════════════════════════════════════════════
# SINGLETON INSTANCES
# ═══════════════════════════════════════════════════════════════

engine = QuantumAssembler()
optimizer = SequenceOptimizer(engine)
bias_optimizer = BiasOptimizer(engine)
compiler = SPPSCompiler()
