import os
import shutil
import tempfile
import pathlib
from typing import Optional, List, Tuple

import ase
import ase.io
from .calculator import OpenMMProteinLigandCalculator
from .scoring import score_structure

from pydantic import BaseModel, Field

import logging
logger = logging.getLogger(__name__)

class DockingJobConfig(BaseModel):
    """
    User-facing configuration for a docking workflow.
    """
    receptor_path: str
    ligands_paths: Optional[List[str]] = None
    ligand_path: Optional[str] = None
    engine: str = "openmm"
    box_center: Optional[List[float]] = None
    box_size: Optional[List[float]] = None
    exhaustiveness: int = 8
    n_poses: int = 10
    protonation_ph: float = 7.4
    optimization_steps: int = 1000
    temperature: float = 300.0
    output_path: str = "docking_run"

    def model_post_init(self, __context=None) -> None:
        # Resolve paths after initialization
        self.receptor_path = str(pathlib.Path(self.receptor_path).resolve())
        
        # Handle single vs multiple ligand specification cleanly
        resolved_ligands = []
        if self.ligands_paths:
            resolved_ligands.extend([str(pathlib.Path(p).resolve()) for p in self.ligands_paths])
            
        if self.ligand_path:
            resolved_ligands.append(str(pathlib.Path(self.ligand_path).resolve()))
            
        if not resolved_ligands:
            raise ValueError("At least one ligand path must be provided. Use 'ligand_path' or 'ligands_paths'.")
            
        self.ligands_paths = resolved_ligands
        # The following attributes are already set by Pydantic's validation and assignment
        # during model initialization, so re-assigning them from non-existent local variables
        # or from self.attribute is redundant.
        # self.engine = self.engine
        # self.box_center = self.box_center
        # self.box_size = self.box_size
        # self.exhaustiveness = self.exhaustiveness
        # self.n_poses = self.n_poses
        # self.protonation_ph = self.protonation_ph
        # self.optimization_steps = self.optimization_steps
        # self.temperature = self.temperature


class DockingResult:
    """
    Standardized output for a docking job.
    """
    def __init__(self, poses: List[ase.Atoms], scores: List[dict]):
        self.poses = poses
        self.scores = scores


class DockingValidator:
    def validate(self, config: DockingJobConfig):
        if not os.path.exists(config.receptor_path):
            raise FileNotFoundError(f"Receptor file not found: {config.receptor_path}")
            
        if not config.ligands_paths:
            raise ValueError("At least one ligand path must be provided.")
            
        for path in config.ligands_paths:
            if not os.path.exists(path):
                raise FileNotFoundError(f"Ligand file not found: {path}")
        
        # We only support openmm for now as a calculator concept
        if config.engine.lower() not in ["openmm", "internal"]:
            raise ValueError(f"Engine {config.engine} not currently supported in this prototype.")


class ReceptorPreparator:
    def prepare(self, path: str, config: DockingJobConfig) -> ase.Atoms:
        # In a real workflow, protonate using pdb2pqr or similar tools
        # For now, just load it
        return ase.io.read(path)


class LigandPreparator:
    def prepare(self, paths: List[str], config: DockingJobConfig) -> List[ase.Atoms]:
        # In a real workflow, generate 3D conformers, assign charges, etc.
        ligands = []
        for p in paths:
            ligands.append(ase.io.read(p))
        return ligands


class DockingEngineInterface:
    def run(self, prepared_receptor: ase.Atoms, prepared_ligands: List[ase.Atoms], config: DockingJobConfig) -> DockingResult:
        raise NotImplementedError()


class InternalOpenMMEngine(DockingEngineInterface):
    """
    A concrete implementation that would normally wire up the GA + OpenMM calculator.
    In this prototype, we just mock the setup to show the architecture.
    """
    def run(self, prepared_receptor: ase.Atoms, prepared_ligands: List[ase.Atoms], config: DockingJobConfig) -> DockingResult:
        import numpy as np
        
        # 1. Combine receptor and ligands into one Atoms object for ASE
        combined = prepared_receptor.copy()
        
        # Track origin entity: 0 for receptor, 1..N for ligands
        entity_ids = [0] * len(combined)
        is_ligand_mask = [0] * len(combined)
        
        for i, lig in enumerate(prepared_ligands):
            combined.extend(lig)
            entity_ids.extend([i + 1] * len(lig))
            is_ligand_mask.extend([1] * len(lig))

        # Mark which atoms belong to the ligands for targeted mutation/crossover
        # This gives operations maximum insight
        combined.set_array('is_ligand', np.array(is_ligand_mask, dtype=int))
        combined.set_array('entity_id', np.array(entity_ids, dtype=int))

        def _get_ase(obj):
            from ase import Atoms
            if isinstance(obj, Atoms): return obj
            if hasattr(obj, 'to_ase'): return obj.to_ase()
            if hasattr(obj, 'get_ase'): return obj.get_ase()
            
            apm = getattr(obj, 'AtomPositionManager', obj)
            pos = getattr(apm, 'atomPositions', None)
            labels = getattr(apm, 'atomLabelsList', None)
            if pos is not None and labels is not None:
                cell = getattr(apm, 'latticeVectors', None)
                return Atoms(symbols=labels, positions=pos, cell=cell)
            return obj

        # 2. Build EZGA Evolutionary Objective for Docking
        from ezga.evaluator.objective import evaluate_objectives
        
        def objective_docking_fitness():
            """Creates a single-objective proxy fitness function for the GA."""
            def compute(dataset):
                from ase import Atoms
                try:
                    structs = list(dataset)
                except TypeError:
                    structs = dataset.structures if hasattr(dataset, 'structures') else [dataset]
                logger.info(f"{dataset.N} {len(structs)} 1231123")

                fits = []
                for s in structs:
                    atoms = _get_ase(s)
                    
                    # Instead of relying on fragile getting from parsed structures, 
                    # we directly use the closure's persistent mask since length is rigid
                    ligand_idx = [idx for idx, is_lig in enumerate(is_ligand_mask) if is_lig == 1]
                    protein_idx = [idx for idx, is_lig in enumerate(is_ligand_mask) if is_lig == 0]
                    lig_types = [atoms[i].symbol for i in ligand_idx]
                    pro_types = [atoms[i].symbol for i in protein_idx]
                    
                    v = score_structure(atoms, ligand_idx, protein_idx, lig_types, pro_types)
                    # Lower is better
                    fitness = v[0] * 5.0 + v[2] * 2.0 + v[3] * 0.5 + v[4] * 1.0
                    fits.append(fitness)
                return np.array(fits, dtype=float)
            return compute

        # 3. Build EZGA Targeted Mutations
        def mutation_rigid_body(std_translation=0.5, std_rotation=0.1):
            """Creates a variation operator targeting only ligand entities."""
            def mutate(structure, **kwargs):
                import copy
                child = copy.deepcopy(structure)
                
                atoms = _get_ase(child)
                pos = atoms.get_positions()
                    
                # Use closure variables for safety
                for lig_id in set(entity_ids):
                    if lig_id != 0:
                        mask = (np.array(entity_ids) == lig_id)
                        center = pos[mask].mean(axis=0)
                        # Translation
                        shift = np.random.uniform(-1.0, 1.0, size=3) * std_translation
                        # Small rotation approximation
                        disp = pos[mask] - center
                        rot = np.eye(3) + np.random.normal(0, std_rotation, (3,3))
                        pos[mask] = center + (disp @ rot) + shift
                        
                # Write back into EZGA structure safely
                if hasattr(child, 'AtomPositionManager'):
                    child.AtomPositionManager.atomPositions = pos
                else:
                    child.set_positions(pos)
                    
                return child
            return mutate

        # 4. OpenMM Calculator for EZGA Framework Compatibility
        from ase.calculators.calculator import Calculator, all_changes
        
        try:
            from openmm import openmm, unit, app
            from ezga.workflows.docking.calculator import OpenMMProteinLigandCalculator
            
            # Proof-of-concept OpenMM Generic System (Lennard-Jones)
            # This demonstrates how to plug openmm without needing complex amber configurations
            # for the arbitrary test elements.
            system = openmm.System()
            for mass in combined.get_masses():
                system.addParticle(mass * unit.dalton)
                
            # Generic repulsive/attractive Lennard-Jones 
            force = openmm.CustomNonbondedForce("4*epsilon*((sigma/r)^12-(sigma/r)^6); sigma=0.5*(sigma1+sigma2); epsilon=sqrt(eps1*eps2)")
            force.addPerParticleParameter("sigma1")
            force.addPerParticleParameter("eps1")
            
            for i in range(len(combined)):
                force.addParticle([0.3 * unit.nanometer, 1.0 * unit.kilojoule_per_mole])
            system.addForce(force)
            
            topology = app.Topology()
            chain = topology.addChain()
            res = topology.addResidue("MOC", chain)
            for sym in combined.get_chemical_symbols():
                el = app.Element.getBySymbol(sym) if sym != 'X' else app.Element.getBySymbol('C')
                topology.addAtom(sym, el, res)
                
            integrator = openmm.VerletIntegrator(0.001 * unit.picoseconds)
            simulation = app.Simulation(topology, system, integrator)
            
            # Boolean mask for the calculator to freeze the receptor natively
            bool_mask = np.array(is_ligand_mask, dtype=bool)
            
            calc = OpenMMProteinLigandCalculator(
                simulation=simulation,
                mask_movable=bool_mask
            )
            logger.info("Successfully mounted OpenMMProteinLigandCalculator.")
            
        except ImportError:
            logger.warning("OpenMM not installed. Using DummyDockingCalculator.")
            class DummyDockingCalculator(Calculator):
                implemented_properties = ["energy", "forces"]
                def calculate(self, atoms=None, properties=("energy", "forces"), system_changes=all_changes):
                    super().calculate(atoms, properties, system_changes)
                    self.results["energy"] = 0.0
                    self.results["forces"] = np.zeros((len(atoms), 3))
            calc = DummyDockingCalculator()

        # 5. Wire the EZGA Engine via Configuration Factory
        import tempfile
        import os
        from ezga.factory import load_config, build_default_engine
        from ezga.simulator.mace_calculator import ase_calculator

        num_generations = max(1, config.optimization_steps // config.n_poses)
        
        output_dir = os.path.abspath(config.output_path)
        os.makedirs(output_dir, exist_ok=True)
        seed_path = os.path.join(output_dir, "docking_seed.xyz")
        
        # Generate the true initial population size
        init_pop = []
        for _ in range(config.n_poses):
            c = combined.copy()
            pos_c = c.get_positions()
            for lig_id in set(entity_ids):
                if lig_id != 0:
                    mask = (np.array(entity_ids) == lig_id)
                    # Minor initial randomization
                    pos_c[mask] += np.random.uniform(-5.0, 5.0, size=3)
            c.set_positions(pos_c)
            init_pop.append(c)
            
        ase.io.write(seed_path, init_pop)
        
        cfg_dict = {
            "output_path": output_dir,
            "max_generations": num_generations,
            "population": {
                "dataset_path": seed_path,
                "filter_duplicates": True,
                "collision_factor": 0.0, # Handled internally by scoring
            },
            "multiobjective": {
                "size": config.n_poses,
            },
            "thermostat": {
                "initial_temperature": config.temperature,
                "temperature_bounds": [0.0, max(config.temperature * 2.0, 1000.0)]
            },
            "evaluator": {
                "objectives_funcs": [objective_docking_fitness()]
            },
            "mutation_funcs": [mutation_rigid_body(std_translation=1.0, std_rotation=0.05)],
            "simulator": {
                "mode": "sampling",
                "calculator": ase_calculator(calculator=calc, steps_max=10, device='cpu')
            }
        }
        
        cfg = load_config(cfg_dict)
        engine = build_default_engine(cfg)
        
        # Execute Full Core Optimization
        final_dataset = engine.run()
        
        # 6. Extract and Format Results
        poses = []
        scores = []
        
        try:
            structs = list(final_dataset)
            structs.sort(key=lambda x: getattr(x, 'scores', [float('inf')])[0])
        except Exception:
            structs = getattr(final_dataset, 'structures', [combined])
            
        for s in structs[:config.n_poses]:
            p = _get_ase(s)
            poses.append(p)
            
            ligand_idx = [idx for idx, is_lig in enumerate(is_ligand_mask) if is_lig == 1]
            protein_idx = [idx for idx, is_lig in enumerate(is_ligand_mask) if is_lig == 0]
            lig_types = [p[i].symbol for i in ligand_idx]
            pro_types = [p[i].symbol for i in protein_idx]
            
            vec = score_structure(p, ligand_idx, protein_idx, lig_types, pro_types)
            fit = vec[0] * 5.0 + vec[2] * 2.0 + vec[3] * 0.5 + vec[4] * 1.0
            
            scores.append({
                "score_vector": vec.tolist(),
                "fitness": fit,
                "clash": vec[0],
                "strain": vec[1],
                "hbond": -vec[2],
                "hydrophobic": -vec[3],
                "burial": -vec[4]
            })

        return DockingResult(poses=poses, scores=scores)


class DockingService:
    """
    A unified facade/service layer for running docking workflows.
    Hides the complexity of receptor/ligand preparation, execution, and postprocessing.
    """
    def __init__(self):
        self.validator = DockingValidator()
        self.receptor_preparator = ReceptorPreparator()
        self.ligand_preparator = LigandPreparator()
        self.engines = {
            "openmm": InternalOpenMMEngine(),
            "internal": InternalOpenMMEngine()
        }

    def run(self, config: DockingJobConfig) -> DockingResult:
        # 1. Validate
        self.validator.validate(config)
        
        # 2. Prepare
        receptor = self.receptor_preparator.prepare(config.receptor_path, config)
        ligands = self.ligand_preparator.prepare(config.ligands_paths, config)
        
        # 3. Select Engine
        engine = self.engines.get(config.engine.lower())
        if not engine:
            raise ValueError(f"Configured engine {config.engine} not found.")
            
        # 4. Docking Execution
        result = engine.run(receptor, ligands, config)
        
        # 5. Output packaging/analytics could go here
        return result
