import warnings

def features_composition(structures):
    """
    Computes a composition-count matrix for a given list of structures.

    Each row i corresponds to 'structures[i]'.
    Each column j corresponds to a unique atom label found across all structures.
    The entry (i, j) is the count of that label in the i-th structure.

    Parameters
    ----------
    structures : list
        List of structures (each having structure.AtomPositionManager.atomLabelsList).
    
    Returns
    -------
    feats : np.ndarray
        2D array of shape (N, D), where N = number of structures,
        D = number of unique atom labels across all structures.
    """
    # You can optionally use the Partition class to gather all unique labels:
    partition = Partition()
    partition.containers = structures

    # 'partition.uniqueAtomLabels' is a set/list of all atom labels across these structures
    atom_labels = partition.uniqueAtomLabels

    # Build a composition vector for each structure, counting how many atoms of each label it has
    feats = np.array([
        [
            np.count_nonzero(structure.AtomPositionManager.atomLabelsList == label)
            for label in atom_labels
        ]
        for structure in structures
    ])

    return feats

def objective_Ef(structures):
    """
    Computes two objective values for each structure:
      1) Anomaly (based on Mahalanobis distance in composition space),
      2) Formation energy (via linear regression on composition).

    Returns a 2D array of shape (N, 2), with columns:
      [ anomaly, formation_energy ] for each structure.

    Parameters
    ----------
    structures : list
        List of structures, each with:
            - structure.AtomPositionManager.E (raw energy)
            - structure.AtomPositionManager.atomLabelsList (atom labels)

    Returns
    -------
    objectives : np.ndarray
        Shape (N, 2). The first column is anomaly, the second is formation energy.
        N = number of structures.
    """
    N = len(structures)
    if N == 0:
        return np.zeros((0, 2))

    # 1) Build composition matrix X (N x D)
    X = features(structures)

    # 2) Extract raw energies into y
    y = np.array([s.AtomPositionManager.E for s in structures])

    # 3) Compute formation energy:
    #    Fit a Ridge model to find chemical potentials, then formation = y - X @ chemical_potentials
    model = Ridge(alpha=1e-5, fit_intercept=False)
    model.fit(X, y)
    chemical_potentials = model.coef_
    formation_energies = y - X.dot(chemical_potentials)

    # 4) Compute anomaly:
    #    We do a Mahalanobis distance for each structure's composition row in X.
    #    "Lower anomaly" means it is 'typical'; "Higher anomaly" means it's more 'outlier'.
    #    For demonstration, we treat all structures as a single distribution.

    # Normalize X (z-score). If a column is constant, it becomes 0.0
    X_norm = zscore(X, axis=0)
    X_norm = np.nan_to_num(X_norm, nan=0.0)
    mean_vec = np.mean(X_norm, axis=0)
    # Small ridge term for stability
    cov_inv = inv(np.cov(X_norm, rowvar=False) + 1e-12 * np.eye(X_norm.shape[1]))

    # Mahalanobis distance => anomaly
    anomalies = np.array([
        mahalanobis(row, mean_vec, cov_inv) for row in X_norm
    ])

    # 5) Pack into a (N,2) result: columns = [ anomaly, formation_energy ]
    objectives = np.column_stack([anomalies, formation_energies])

    return objectives

def validate_evolution_parameters(params: dict = None) -> dict:
    """
    Validates and corrects evolutionary parameters for the structure exploration workflow.
    
    This function ensures that the input parameters dictionary contains all necessary keys with 
    appropriate types. Missing keys are populated with default values and warnings are issued 
    when corrections occur. If a parameter's type does not match the expected type, an AssertionError is raised.
    
    Parameters
    ----------
    params : dict, optional
        Dictionary containing parameter overrides. Expected keys and their default values are:
            - 'max_generations' (int): Maximum number of generations to run; default is 100.
            - 'min_size_for_filter' (int): Minimum dataset size for filtering; default is 10.
            - 'mutation_rate_params' (dict): Dictionary with keys:
                * 'min_mutation_rate' (int): Default is 1.
                * 'initial_mutation_rate' (float): Default is 100.0.
                * 'max_prob' (float): Default is 0.9.
                * 'min_prob' (float): Default is 0.1.
                * 'alpha' (float): Default is 0.01.
                * 'use_magnitude_scaling' (bool): Default is True.
            - 'convergence_params' (dict): Dictionary with keys:
                * 'objective_threshold' (float): Default is 0.01.
                * 'feature_threshold' (float): Default is 0.01.
                * 'stall_threshold' (int): Default is 5.
                * 'information_driven' (bool) : Default is False.
            - 'collision_factor' (float): Factor for collision filtering; default is 0.4.
            - 'md_params' (dict): Parameters for MD simulations; default is {}.
            - 'foreigners' (int): Placeholder for foreigners count; default is 0.
            - 'dataset_path' (str): Path to the dataset file; default is 'dataset.xyz'.
            - 'template_path' (str): Path to the template file; default is 'template.xyz'.
            - 'output_path' (str): Directory to save outputs; default is '.'.
            - 'objective_funcs', 'features_funcs', 'mutation_funcs', 'crossover_funcs',
              'DoE', 'BO', 'thermostat', 'physical_model_func': Workflow functions; default is None.
            - 'multiobjective_params' (dict): Dictionary with keys:
                * 'weights' (None, list, or tuple): Default is None.
                * 'repulsion_weight' (float): Default is 1.0.
                * 'size' (int): Default is 20.
                * 'repulsion_mode' (str): Default is 'avg'.
                * 'metric' (str): Default is 'euclidean'.
                * 'random_seed' (int): Default is 73.
    
    Returns
    -------
    dict
        A validated and corrected dictionary of evolutionary parameters.
    
    Raises
    ------
    AssertionError
        If any parameter does not match the expected type.
    """
    # Define default parameters

    defaults = {
        'max_generations': 100,
        'min_size_for_filter': 10,
        'mutation_rate_params': {
            'min_mutation_rate': 1,
            'initial_mutation_rate': 100.0,
            'max_prob': 0.9,
            'min_prob': 0.1,
            'alpha':0.01,
            'use_magnitude_scaling':True,
            'crossover_probability':0.1,  
        },
        'convergence_params': {
            'objective_threshold': 0.01,
            'feature_threshold': 0.01,
            'stall_threshold': 5,  
            'information_driven':False
        },
        'multiobjective_params':{
            'weights': None,
            'repulsion_weight': 1.0,
            'repetition_penalty': False,
            'size': 20,
            'repulsion_mode': 'avg',
            'metric': 'euclidean',
            'random_seed': 73,
            'steepness': 10.0,
            'max_count': 100,
            'cooling_rate': 0.1,
        },
        'information_metric':{
            'auto_percentile': 15,   
            'metric': 'novelty',   
            'components': 5,   
            'r_cut': 5.0,   
            'n_max': 3,   
            'l_max': 3,   
            'sigma': 0.5,   
            'max_clusters': 10,   
        },
        'collision_factor': 0.4,
        'md_params': {},
        'foreigners': 0,
        'dataset_path': 'dataset.xyz',
        'template_path': 'template.xyz',
        'output_path': '.',
        'objective_funcs': objective_Ef,
        'features_funcs': features_composition,
        'mutation_funcs': None,
        'crossover_funcs': None,
        'DoE': None,
        'BO': None,
        'thermostat': None,
        'physical_model_func': None,
    }

    # If no parameters provided, use an empty dictionary
    if params is None:
        params = {}

    # Create a copy of defaults to update and return
    validated_params = defaults.copy()

    # Update top-level parameters with provided values, if any
    top_level_keys = [
        'max_generations', 'min_size_for_filter', 'collision_factor', 
        'md_params', 'foreigners', 'dataset_path', 'template_path', 'output_path',
        'objective_funcs', 'features_funcs', 'mutation_funcs', 'crossover_funcs',
        'DoE', 'BO', 'thermostat', 'physical_model_func', 
    ]
    for key in top_level_keys:
        if key in params:
            validated_params[key] = params[key]

    # Validate top-level parameters
    assert isinstance(validated_params['max_generations'], int), (
        f"'max_generations' must be an int, got {type(validated_params['max_generations']).__name__}"
    )
    assert isinstance(validated_params['min_size_for_filter'], int), (
        f"'min_size_for_filter' must be an int, got {type(validated_params['min_size_for_filter']).__name__}"
    )
    assert isinstance(validated_params['collision_factor'], float), (
        f"'collision_factor' must be a float, got {type(validated_params['collision_factor']).__name__}"
    )
    assert isinstance(validated_params['md_params'], dict), (
        f"'md_params' must be a dict, got {type(validated_params['md_params']).__name__}"
    )
    if isinstance(validated_params['foreigners'], float):
        warnings.warn("Parameter 'foreigners' provided as float. Converting to int.")
        validated_params['foreigners'] = int(validated_params['foreigners'])
    assert isinstance(validated_params['foreigners'], int), (
        f"'foreigners' must be an int, got {type(validated_params['foreigners']).__name__}"
    )
    for path_key in ['dataset_path', 'template_path', 'output_path']:
        assert isinstance(validated_params[path_key], str) and validated_params[path_key], (
            f"'{path_key}' must be a non-empty string."
        )

    # Validate mutation_rate_params
    mutation_defaults = defaults['mutation_rate_params']
    validated_params['mutation_rate_params'] = params.get('mutation_rate_params', mutation_defaults)
    assert isinstance(validated_params['mutation_rate_params'], dict), (
        f"'mutation_rate_params' must be a dict, got {type(validated_params['mutation_rate_params']).__name__}"
    )
    expected_mutation_keys = {
        'min_mutation_rate': int,
        'initial_mutation_rate': float,
        'max_prob': float,
        'min_prob': float,
        'alpha':float,
        'use_magnitude_scaling':bool,
        'crossover_probability':float
        ,  
    }
    for key, expected_type in expected_mutation_keys.items():
        if key not in validated_params['mutation_rate_params']:
            warnings.warn(f"Key '{key}' missing in 'mutation_rate_params'. Using default value: {mutation_defaults[key]}")
            validated_params['mutation_rate_params'][key] = mutation_defaults[key]
        else:
            assert isinstance(validated_params['mutation_rate_params'][key], expected_type), (
                f"Key '{key}' in 'mutation_rate_params' must be {expected_type.__name__}, "
                f"got {type(validated_params['mutation_rate_params'][key]).__name__}"
            )

    # Validate convergence_params
    conv_defaults = defaults['convergence_params']
    validated_params['convergence_params'] = params.get('convergence_params', conv_defaults)
    assert isinstance(validated_params['convergence_params'], dict), (
        f"'convergence_params' must be a dict, got {type(validated_params['convergence_params']).__name__}"
    )
    expected_conv_keys = {
        'objective_threshold': (float, int),
        'feature_threshold': (float, int),
        'stall_threshold': int,
        'information_driven': bool,
    }
    for key, expected_type in expected_conv_keys.items():
        if key not in validated_params['convergence_params']:
            warnings.warn(f"Key '{key}' missing in 'convergence_params'. Using default value: {conv_defaults[key]}")
            validated_params['convergence_params'][key] = conv_defaults[key]
        else:
            assert isinstance(validated_params['convergence_params'][key], expected_type), (
                f"Key '{key}' in 'convergence_params' must be {expected_type.__name__}, "
                f"got {type(validated_params['convergence_params'][key]).__name__}"
            )

    # ----- Validate multiobjective_params -----
    multi_defaults = defaults['multiobjective_params']
    validated_params['multiobjective_params'] = params.get('multiobjective_params', multi_defaults)
    assert isinstance(validated_params['multiobjective_params'], dict), (
        f"'multiobjective_params' must be a dict, got {type(validated_params['multiobjective_params']).__name__}"
    )

    expected_multi_keys = {
        'weights': (type(None), list, tuple),
        'repulsion_weight': (float, int),
        'repetition_penalty': bool,
        'size': int,
        'repulsion_mode': str,
        'metric': str,
        'random_seed': int,
        'steepness': (float, int),
        'max_count': (float, int),
        'cooling_rate': (float, int),
    }
    for key, expected_type in expected_multi_keys.items():
        if key not in validated_params['multiobjective_params']:
            warnings.warn(f"Key '{key}' missing in 'multiobjective_params'. Using default value: {multi_defaults[key]}")
            validated_params['multiobjective_params'][key] = multi_defaults[key]
        else:
            if isinstance(expected_type, tuple):
                if not isinstance(validated_params['multiobjective_params'][key], expected_type):
                    raise AssertionError(
                        f"Key '{key}' in 'multiobjective_params' must be one of types {[t.__name__ for t in expected_type]}, "
                        f"got {type(validated_params['multiobjective_params'][key]).__name__}"
                    )
            else:
                assert isinstance(validated_params['multiobjective_params'][key], expected_type), (
                    f"Key '{key}' in 'multiobjective_params' must be {expected_type.__name__}, "
                    f"got {type(validated_params['multiobjective_params'][key]).__name__}"
                )

    # Retrieve size from multiobjective_params
    size = validated_params['multiobjective_params']['size']


    # ----- Validate information_metric -----
    information_defaults = defaults['information_metric']
    validated_params['information_metric'] = params.get('information_metric', information_defaults)
    assert isinstance(validated_params['information_metric'], dict), (
        f"'information_metric' must be a dict, got {type(validated_params['information_metric']).__name__}"
    )

    expected_info_keys = {
        'auto_percentile': (float, int),
        'metric': str,
        'components': int,    
        'r_cut': (float, int),
        'n_max': int,
        'l_max': int,
        'sigma': (float, int),
        'max_clusters': int,
    }
    for key, expected_type in expected_info_keys.items():
        if key not in validated_params['information_metric']:
            warnings.warn(f"Key '{key}' missing in 'information_metric'. Using default value: {information_defaults[key]}")
            validated_params['information_metric'][key] = information_defaults[key]
        else:
            if isinstance(expected_type, tuple):
                if not isinstance(validated_params['information_metric'][key], expected_type):
                    raise AssertionError(
                        f"Key '{key}' in 'information_metric' must be one of types {[t.__name__ for t in expected_type]}, "
                        f"got {type(validated_params['information_metric'][key]).__name__}"
                    )
            else:
                assert isinstance(validated_params['information_metric'][key], expected_type), (
                    f"Key '{key}' in 'information_metric' must be {expected_type.__name__}, "
                    f"got {type(validated_params['information_metric'][key]).__name__}"
                )

    # Check if min_size_for_filter is less than size and update if necessary
    if validated_params['min_size_for_filter'] < size:
        warnings.warn(
            f"'min_size_for_filter' ({validated_params['min_size_for_filter']}) is less than "
            f"'size' ({size}). Setting 'min_size_for_filter' to {size}."
        )
        validated_params['min_size_for_filter'] = size

    # --- Additional check for unrecognized hyperparameters ---

    # Top-level: warn about any extra keys in params that are not defined in defaults.
    for key in params:
        if key not in defaults:
            warnings.warn(f"Unrecognized hyperparameter '{key}' provided. It will be ignored.")

    # For nested dictionaries: check for extra keys in 'mutation_rate_params'
    if 'mutation_rate_params' in params and isinstance(params['mutation_rate_params'], dict):
        for key in params['mutation_rate_params']:
            if key not in expected_mutation_keys:
                warnings.warn(f"Unrecognized hyperparameter 'mutation_rate_params->{key}' provided. It will be ignored.")

    # For nested dictionaries: check for extra keys in 'convergence_params'
    if 'convergence_params' in params and isinstance(params['convergence_params'], dict):
        for key in params['convergence_params']:
            if key not in expected_conv_keys:
                warnings.warn(f"Unrecognized hyperparameter 'convergence_params->{key}' provided. It will be ignored.")

    # For nested dictionaries: check for extra keys in 'multiobjective_params'
    if 'multiobjective_params' in params and isinstance(params['multiobjective_params'], dict):
        for key in params['multiobjective_params']:
            if key not in expected_multi_keys:
                warnings.warn(f"Unrecognized hyperparameter 'multiobjective_params->{key}' provided. It will be ignored.")

    # For nested dictionaries: check for extra keys in 'multiobjective_params'
    if 'information_metric' in params and isinstance(params['information_metric'], dict):
        for key in params['information_metric']:
            if key not in expected_info_keys:
                warnings.warn(f"Unrecognized hyperparameter 'information_metric->{key}' provided. It will be ignored.")

    return validated_params