# -*- coding: utf-8 -*-
"""
F5 - Simulacion rigurosa del dodecamero 1BNA con OpenMM. INTERRUMPIBLE.

Primera ejecucion: construye el sistema (1BNA -> +H -> agua TIP3P 1 nm
+ NaCl 150 mM), minimiza, equilibra y produce. Ademas de la trayectoria,
congela la DEFINICION del sistema (f5_system.xml, f5_sistema.pdb,
f5_ref.npz) para que un checkpoint pueda aplicarse mas tarde.

Se puede cortar cuando sea (Ctrl+C, apagar, lo que toque): el checkpoint
se guarda tras CADA fotograma (a 50 ns/dia, ~1,5 min de GPU en riesgo).

Reanudar:   python f5_simular.py --resume --ns 100
(continua hasta completar los ns pedidos; se puede repetir mil veces)

Salidas en trayectorias/:
  f5_visor.bin   trayectoria para trayectoria.html ('CELL1' + uint32
                 natoms + codigos de elemento + por fotograma: float32
                 t_ps, float32 rmsd_A, natoms*3 float32 en angstroms)
  f5_full.dcd    trayectoria completa estandar (VMD/ChimeraX)
  f5_log.csv     termodinamica y velocidad
  f5_sistema.pdb + f5_system.xml + f5_ref.npz   definicion congelada
  f5.chk         checkpoint (posiciones, velocidades, RNG...)

Uso:  python f5_simular.py --test            prueba corta (0,2 ns)
      python f5_simular.py --ns 100          la de verdad (~2 dias GPU)
      python f5_simular.py --resume --ns 100 retomar donde quedo
"""
import argparse, os, struct, sys, time
import numpy as np
from openmm.app import (PDBFile, Modeller, ForceField, Simulation,
                        PME, HBonds, DCDReporter, StateDataReporter)
from openmm import (LangevinMiddleIntegrator, MonteCarloBarostat,
                    Platform, XmlSerializer, System, Context, VerletIntegrator)
from openmm.unit import kelvin, picosecond, picoseconds, nanometer, molar, bar, angstrom

ap = argparse.ArgumentParser()
ap.add_argument('--ns', type=float, default=1.0, help='ns TOTALES de produccion (objetivo)')
ap.add_argument('--test', action='store_true', help='prueba corta (0,2 ns, fotograma cada 2 ps)')
ap.add_argument('--stride_ps', type=float, default=50.0, help='ps entre fotogramas del visor')
ap.add_argument('--resume', action='store_true', help='reanudar desde el checkpoint')
ap.add_argument('--camara', type=str, default=None, metavar='NOMBRE',
                help='camara de alta velocidad: tramo corto desde el checkpoint a un fichero aparte (f5_cam_NOMBRE.bin)')
ap.add_argument('--ps', type=float, default=4.0, help='duracion del tramo de camara (ps)')
args = ap.parse_args()
if args.test:
    args.ns, args.stride_ps = 0.2, 2.0
if args.camara:
    args.resume = True          # la camara siempre parte del checkpoint congelado

AQUI = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(AQUI, 'trayectorias')
os.makedirs(OUT, exist_ok=True)
def ruta(n): return os.path.join(OUT, n)

def log(msg):
    print(f'[{time.strftime("%H:%M:%S")}] {msg}', flush=True)

def plataforma():
    """CUDA si existe; si no, OpenCL ELIGIENDO la GPU discreta: el defecto
    de OpenMM puede quedarse con la Intel integrada (comprobado: 52,8
    ns/dia contra 125,6 de la GTX 1650 con el sistema F5)."""
    try:
        return Platform.getPlatformByName('CUDA'), {'Precision': 'mixed'}
    except Exception:
        pass
    try:
        plat = Platform.getPlatformByName('OpenCL')
        mejor = None
        for pi in range(4):
            for di in range(4):
                try:
                    dummy = System(); dummy.addParticle(1.0)
                    integ = VerletIntegrator(0.001)
                    ctx = Context(dummy, integ, plat,
                                  {'OpenCLPlatformIndex': str(pi), 'DeviceIndex': str(di)})
                    nom = plat.getPropertyValue(ctx, 'DeviceName')
                    del ctx, integ
                    puntos = 2 if any(s in nom for s in ('NVIDIA', 'GeForce', 'RTX', 'Radeon', 'RX')) else 1
                    if mejor is None or puntos > mejor[0]:
                        mejor = (puntos, pi, di, nom)
                except Exception:
                    pass
        if mejor:
            log(f'GPU elegida: {mejor[3]} (OpenCL plataforma {mejor[1]}, dispositivo {mejor[2]})')
            return plat, {'OpenCLPlatformIndex': str(mejor[1]),
                          'DeviceIndex': str(mejor[2]), 'Precision': 'mixed'}
        return plat, {'Precision': 'mixed'}
    except Exception:
        pass
    return Platform.getPlatformByName('CPU'), {}

DT_PS = 0.002          # 2 fs con SHAKE (HBonds)
CAB_MAGIA = b'CELL1'

# ════════════════════════════════════════════════════════════════
if args.resume:
    # ---------- REANUDAR: cargar la definicion congelada ----------
    for f in ('f5_sistema.pdb', 'f5_system.xml', 'f5_ref.npz', 'f5.chk', 'f5_visor.bin'):
        if not os.path.exists(ruta(f)):
            sys.exit(f'falta trayectorias/{f}: no hay nada que reanudar (ejecuta primero sin --resume)')
    log('reanudando: cargando definicion congelada del sistema')
    ref = np.load(ruta('f5_ref.npz'))
    sel_visor, codigos = ref['sel_visor'], ref['codigos']
    sel_rmsd, ref_rmsd = ref['sel_rmsd'], ref['ref_rmsd']
    stride_guardado = float(ref['stride_ps'])
    if args.camara is None:
        if abs(stride_guardado - args.stride_ps) > 1e-9 and not args.test:
            log(f'aviso: el fichero usa {stride_guardado} ps/fotograma; se respeta ese valor')
        args.stride_ps = stride_guardado   # la pelicula larga no mezcla strides; la camara usa el suyo
    pdb_sis = PDBFile(ruta('f5_sistema.pdb'))
    with open(ruta('f5_system.xml')) as f:
        system = XmlSerializer.deserialize(f.read())
    integrator = LangevinMiddleIntegrator(310 * kelvin, 1.0 / picosecond, DT_PS * picoseconds)
    plat, props = plataforma()
    sim = Simulation(pdb_sis.topology, system, integrator, plat, props)
    sim.loadCheckpoint(ruta('f5.chk'))
    t_chk = sim.context.getState().getTime().value_in_unit(picoseconds)
    sim.currentStep = int(round(t_chk / DT_PS))
    log(f'plataforma: {plat.getName()} · checkpoint en t={t_chk:.1f} ps')
    # fotogramas ya escritos (truncando cualquier resto a medio escribir)
    cab = len(CAB_MAGIA) + 4 + len(codigos)
    por_frame = 8 + len(sel_visor) * 12
    tam = os.path.getsize(ruta('f5_visor.bin'))
    frames_disco = max(0, (tam - cab) // por_frame)
    # el checkpoint (estado fisico) manda: si el .bin tiene UN fotograma mas
    # que el checkpoint (kill justo entre escribir el frame y guardar el chk),
    # ese frame quedaria huerfano — se descarta para que .bin y checkpoint
    # avancen exactamente a la par. ps_eq=100 en produccion (10 en --test).
    ps_eq = 10 if args.test else 100
    frames_chk = int(round((t_chk - ps_eq) / stride_guardado))
    frames_hechos = max(0, min(frames_disco, frames_chk))
    tam_bueno = cab + frames_hechos * por_frame
    if tam != tam_bueno:
        log(f'reconciliando con el checkpoint: {frames_disco} en .bin, {frames_chk} en chk '
            f'-> se conservan {frames_hechos} ({tam - tam_bueno} bytes descartados)')
        with open(ruta('f5_visor.bin'), 'r+b') as f:
            f.truncate(tam_bueno)
    log(f'fotogramas ya en disco: {frames_hechos}')
else:
    # ---------- PRIMERA VEZ: construir el sistema ----------
    if os.path.exists(ruta('f5.chk')):
        log('aviso: habia una simulacion previa; se empieza de cero y se sobreescribe (usa --resume para continuarla)')
    log('leyendo data/1BNA.pdb (cristal de Dickerson, 1981)')
    pdb = PDBFile(os.path.join(AQUI, 'data', '1BNA.pdb'))
    ref_cristal = {}
    for atom, xyz in zip(pdb.topology.atoms(), pdb.positions):
        if atom.residue.name in ('DA', 'DT', 'DG', 'DC') and atom.element.symbol != 'H':
            ref_cristal[(atom.residue.chain.id, atom.residue.id, atom.name)] = \
                np.array([xyz.x, xyz.y, xyz.z]) * 10.0
    ff = ForceField('amber14-all.xml', 'amber14/tip3p.xml')      # ADN: OL15
    mod = Modeller(pdb.topology, pdb.positions)
    mod.deleteWater()
    mod.addHydrogens(ff)
    log('solvatando: 1 nm de margen, NaCl 150 mM + neutralizacion')
    mod.addSolvent(ff, padding=1.0 * nanometer, ionicStrength=0.15 * molar)
    log(f'sistema: {mod.topology.getNumAtoms()} atomos')
    system = ff.createSystem(mod.topology, nonbondedMethod=PME,
                             nonbondedCutoff=1.0 * nanometer, constraints=HBonds)
    system.addForce(MonteCarloBarostat(1 * bar, 310 * kelvin))
    with open(ruta('f5_system.xml'), 'w') as f:
        f.write(XmlSerializer.serialize(system))
    integrator = LangevinMiddleIntegrator(310 * kelvin, 1.0 / picosecond, DT_PS * picoseconds)
    plat, props = plataforma()
    sim = Simulation(mod.topology, system, integrator, plat, props)
    sim.context.setPositions(mod.positions)
    log(f'plataforma: {plat.getName()}')

    CODIGO = {'C': 0, 'N': 1, 'O': 2, 'P': 3, 'H': 4, 'Na': 5, 'Cl': 6}
    sel_visor, codigos, sel_rmsd, ref_rmsd = [], [], [], []
    for atom in mod.topology.atoms():
        res = atom.residue
        if res.name in ('DA', 'DT', 'DG', 'DC'):
            sel_visor.append(atom.index)
            codigos.append(CODIGO[atom.element.symbol])
            if atom.element.symbol != 'H':
                clave = (res.chain.id, res.id, atom.name)
                if clave in ref_cristal:
                    sel_rmsd.append(atom.index)
                    ref_rmsd.append(ref_cristal[clave])
        elif res.name in ('NA', 'CL'):
            sel_visor.append(atom.index)
            codigos.append(CODIGO['Na'] if res.name == 'NA' else CODIGO['Cl'])
    sel_visor = np.array(sel_visor); codigos = np.array(codigos, dtype=np.uint8)
    sel_rmsd = np.array(sel_rmsd)
    ref_rmsd = np.array(ref_rmsd); ref_rmsd -= ref_rmsd.mean(axis=0)
    np.savez(ruta('f5_ref.npz'), sel_visor=sel_visor, codigos=codigos,
             sel_rmsd=sel_rmsd, ref_rmsd=ref_rmsd, stride_ps=args.stride_ps)
    log(f'visor: {len(sel_visor)} atomos · RMSD sobre {len(sel_rmsd)} pesados vs cristal')

    log('minimizando energia...')
    sim.minimizeEnergy()
    estado = sim.context.getState(positions=True)
    with open(ruta('f5_sistema.pdb'), 'w') as f:
        PDBFile.writeFile(mod.topology, estado.getPositions(), f)
    sim.context.setVelocitiesToTemperature(310 * kelvin)
    ps_eq = 10 if args.test else 100
    log(f'equilibrando {ps_eq} ps a 310 K y 1 bar...')
    sim.step(int(ps_eq / DT_PS))
    with open(ruta('f5_visor.bin'), 'wb') as f:
        f.write(CAB_MAGIA)
        f.write(struct.pack('<I', len(sel_visor)))
        f.write(codigos.tobytes())
    frames_hechos = 0

# ════════════════════════════════════════════════════════════════
def rmsd_vs_cristal(pos_A):
    P = pos_A[sel_rmsd] - pos_A[sel_rmsd].mean(axis=0)
    H = P.T @ ref_rmsd
    U, S, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(Vt.T @ U.T))
    R = Vt.T @ np.diag([1, 1, d]) @ U.T
    return float(np.sqrt(np.mean(np.sum((P @ R.T - ref_rmsd) ** 2, axis=1))))

# ---------- modo camara de alta velocidad ----------
# Tramo corto con fotograma cada pocos femtosegundos, a fichero aparte:
# las vibraciones son estadisticamente iguales siempre, no hace falta
# grabarlas 100 ns. No toca f5_visor.bin; el checkpoint avanza esos ps
# (un desvio despreciable en la linea temporal de la pelicula larga).
if args.camara:
    pasos_cam = max(1, int(round(args.stride_ps / DT_PS)))
    stride_real = pasos_cam * DT_PS
    nfr_cam = max(1, int(round(args.ps / stride_real)))
    fichero = ruta(f'f5_cam_{args.camara}.bin')
    tam_est = (8 + len(sel_visor) * 12) * nfr_cam / 1e6
    log(f'camara "{args.camara}": {args.ps} ps · fotograma cada {stride_real * 1000:.1f} fs '
        f'· {nfr_cam} fotogramas · ~{tam_est:.1f} MB -> {os.path.basename(fichero)}')
    with open(fichero, 'wb') as f:
        f.write(CAB_MAGIA)
        f.write(struct.pack('<I', len(sel_visor)))
        f.write(np.asarray(codigos, dtype=np.uint8).tobytes())
    t0, ultimo_chk = time.time(), time.time()
    try:
        for i in range(nfr_cam):
            sim.step(pasos_cam)
            st = sim.context.getState(positions=True)
            pos = np.array(st.getPositions().value_in_unit(angstrom))
            r = rmsd_vs_cristal(pos)
            with open(fichero, 'ab') as f:
                f.write(struct.pack('<ff', (i + 1) * stride_real, r))
                f.write(pos[sel_visor].astype('<f4').tobytes())
            # aqui el checkpoint por fotograma seria carisimo: cada 15 s basta
            if time.time() - ultimo_chk > 15:
                sim.saveCheckpoint(ruta('f5.chk'))
                ultimo_chk = time.time()
            if (i + 1) % max(1, nfr_cam // 10) == 0:
                log(f'  {i + 1}/{nfr_cam} fotogramas · {(time.time() - t0):.0f} s')
    except KeyboardInterrupt:
        log('camara interrumpida: el tramo grabado hasta aqui es valido.')
    sim.saveCheckpoint(ruta('f5.chk'))
    log(f'camara completada. Ver con: trayectoria.html?traj={os.path.basename(fichero)}')
    sys.exit(0)

pasos_stride = int(round(args.stride_ps / DT_PS))
n_frames = max(1, int(round(args.ns * 1000 / args.stride_ps)))
if frames_hechos >= n_frames:
    sys.exit(f'ya hay {frames_hechos} fotogramas (objetivo {n_frames}): sube --ns si quieres mas')
log(f'produccion: objetivo {args.ns} ns ({n_frames} fotogramas de {args.stride_ps} ps) · '
    f'hechos {frames_hechos} · quedan {n_frames - frames_hechos}')

sim.reporters.append(DCDReporter(ruta('f5_full.dcd'), max(pasos_stride, 50000),
                                 append=args.resume and os.path.exists(ruta('f5_full.dcd'))))
sim.reporters.append(StateDataReporter(open(ruta('f5_log.csv'), 'a' if args.resume else 'w'),
                     5000, step=True, time=True, temperature=True,
                     potentialEnergy=True, density=True, speed=True, separator=','))

t0, hechos_aqui = time.time(), 0
try:
    for i in range(frames_hechos, n_frames):
        sim.step(pasos_stride)
        st = sim.context.getState(positions=True)
        pos = np.array(st.getPositions().value_in_unit(angstrom))
        r = rmsd_vs_cristal(pos)
        t_ps = (i + 1) * args.stride_ps
        with open(ruta('f5_visor.bin'), 'ab') as f:
            f.write(struct.pack('<ff', t_ps, r))
            f.write(pos[sel_visor].astype('<f4').tobytes())
        sim.saveCheckpoint(ruta('f5.chk'))
        hechos_aqui += 1
        if hechos_aqui == 1 or i % max(1, n_frames // 40) == 0 or i == n_frames - 1:
            ns_dia = (hechos_aqui * args.stride_ps / 1000) / ((time.time() - t0) / 86400)
            log(f'  t={t_ps:8.1f} ps · RMSD={r:5.2f} A · {ns_dia:6.1f} ns/dia · '
                f'{i + 1}/{n_frames} fotogramas')
except KeyboardInterrupt:
    log(f'interrumpido a peticion: {frames_hechos + hechos_aqui}/{n_frames} fotogramas en disco.')
    log('reanudable con:  python f5_simular.py --resume --ns ' + str(args.ns))
    sys.exit(0)

log('produccion completada.')