import copy
import random
import math
import numpy as np
from typing import List, Any, Tuple

def make_gibbs_exchange_crossover(gibbs_cfg: Any, logger: Any = None):
    """
    Returns a standard crossover function (P1, P2) -> (C1, C2) that performs 
    Gibbs ensemble moves.
    """
    def gibbs_crossover(parent1: Any, parent2: Any) -> Tuple[Any, Any]:
        child1 = copy.deepcopy(parent1)
        child2 = copy.deepcopy(parent2)
        
        apm1 = getattr(child1, "AtomPositionManager", None)
        apm2 = getattr(child2, "AtomPositionManager", None)
        
        if not apm1 or not apm2:
            return child1, child2

        # Choose move type based on weights
        vol_weight = gibbs_cfg.volume_weight
        ptr_weight = gibbs_cfg.transfer_weight
        
        move_type = "none"
        
        if random.random() < (vol_weight / (vol_weight + ptr_weight)):
            # --- VOLUME EXCHANGE ---
            if gibbs_cfg.allow_volume_exchange:
                move_type = "volume"
                V1_old = apm1.get_volume()
                V2_old = apm2.get_volume()
                V_total = V1_old + V2_old
                
                # Linear volume jump (symmetric in V space)
                delta_v = random.uniform(-1.0, 1.0) * (0.05 * V_total)
                V1_new = V1_old + delta_v
                
                if 0.01 * V_total < V1_new < 0.99 * V_total:
                    V2_new = V_total - V1_new
                    scale1 = (V1_new / V1_old) ** (1/3)
                    scale2 = (V2_new / V2_old) ** (1/3)
                    apm1.set_latticeVectors(apm1.latticeVectors * scale1)
                    apm2.set_latticeVectors(apm2.latticeVectors * scale2)
        else:
            # --- PARTICLE SWAP ---
            move_type = "particle"
            
            # 1->2 or 2->1?
            if random.random() < 0.5:
                s_apm, d_apm, s_w, d_w = apm1, apm2, child1, child2
            else:
                s_apm, d_apm, s_w, d_w = apm2, apm1, child2, child1

            if len(s_apm.atomPositions) > 0:
                atom_idx = random.randint(0, len(s_apm.atomPositions) - 1)
                species = s_apm.atomLabelsList[atom_idx]
                
                # Remove from source
                s_apm.atomLabelsList = np.delete(s_apm.atomLabelsList, atom_idx)
                s_apm.atomPositions = np.delete(s_apm.atomPositions, atom_idx, axis=0)
                
                # Add to destination (random position)
                frac = [random.random() for _ in range(3)]
                new_pos = np.dot(frac, d_apm.latticeVectors)
                d_apm.atomLabelsList = np.append(d_apm.atomLabelsList, species)
                if len(d_apm.atomPositions) == 0:
                    d_apm.atomPositions = np.array([new_pos])
                else:
                    d_apm.atomPositions = np.vstack([d_apm.atomPositions, new_pos])
                
                # Tag species for transition kernel
                for c in [child1, child2]:
                    if not hasattr(c, "info"): c.info = {}
                    c.info["gibbs_species"] = species

        # Tag children for the transition kernel
        pair_id = random.randint(0, 1000000)
        for c in [child1, child2]:
            tags = {
                "transformation": "gibbs_exchange",
                "gibbs_move_type": move_type,
                "gibbs_pair_id": pair_id
            }
            apm = getattr(c, "AtomPositionManager", None)
            if apm is not None:
                apm.metadata.update(tags)
            
            if not hasattr(c, "info"): c.info = {}
            c.info.update(tags)
            
        return child1, child2

    return gibbs_crossover
