"""Load and manage Disco103 pre-trained weights.

This module handles loading the Disco103 meta-parameters from the npz file
provided in the DiscoRL repository and integrating them with DiscoRL agents.
"""

import os
from typing import Any, Dict, Tuple
import numpy as np
import jax
import jax.numpy as jnp


def load_disco103_weights(disco_rl_path: str = None) -> Dict[str, Any]:
    """Load Disco103 pre-trained weights.
    
    Args:
        disco_rl_path: path to disco_rl repo root. If None, search from cwd.
    
    Returns:
        dict with keys like 'meta_params' or structure matching the npz file
    """
    if disco_rl_path is None:
        # Try to find disco_rl in standard locations
        possible_paths = [
            'disco_rl/disco_rl/update_rules/weights/disco_103.npz',
            '../disco_rl/disco_rl/update_rules/weights/disco_103.npz',
            '../../disco_rl/disco_rl/update_rules/weights/disco_103.npz',
        ]
        npz_path = None
        for p in possible_paths:
            if os.path.exists(p):
                npz_path = p
                break
        if npz_path is None:
            raise FileNotFoundError(
                'Could not find disco_103.npz. Please provide disco_rl_path '
                'or ensure disco_rl/ is accessible.'
            )
    else:
        npz_path = os.path.join(
            disco_rl_path, 'disco_rl/update_rules/weights/disco_103.npz'
        )
    
    if not os.path.exists(npz_path):
        raise FileNotFoundError(f'disco_103.npz not found at {npz_path}')
    
    # Load the npz file
    data = np.load(npz_path, allow_pickle=True)
    
    # Convert to dictionary; npz files can be accessed as dict-like
    weights = {}
    for key in data.files:
        item = data[key]
        # Some items might be numpy object arrays (e.g., nested structures)
        # Try to convert to jax arrays where possible
        if isinstance(item, np.ndarray):
            weights[key] = jnp.asarray(item)
        else:
            weights[key] = item
    
    print(f'Loaded Disco103 weights from {npz_path}')
    print(f'  Keys: {list(weights.keys())}')
    for key, val in weights.items():
        if hasattr(val, 'shape'):
            print(f'    {key}: shape={val.shape}, dtype={val.dtype}')
        else:
            print(f'    {key}: {type(val)}')
    
    return weights


def unflatten_disco_weights(flat_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Convert flat npz weight dict to nested structure expected by DiscoRL.
    
    The exact structure depends on how the weights were saved. This is a
    placeholder; you may need to adjust based on the actual npz structure.
    
    For now, we assume the npz contains the meta_params directly or under
    a 'meta_params' key.
    """
    # If there's a 'meta_params' key, use it; otherwise assume flat_dict IS the params
    if 'meta_params' in flat_dict:
        meta_params = flat_dict['meta_params']
    else:
        # Try to reconstruct nested structure from flat keys
        # This is environment-specific; adjust as needed
        meta_params = flat_dict
    
    return meta_params


def merge_weights_with_agent(
    agent_meta_state: Dict[str, Any],
    disco_weights: Dict[str, Any],
) -> Dict[str, Any]:
    """Merge loaded Disco103 weights into agent's meta_state.
    
    This updates the meta_state's rnn_state and other components with
    pre-trained weights if available.
    
    Args:
        agent_meta_state: the agent's initial meta_state dict
        disco_weights: loaded weights dict
    
    Returns:
        updated agent_meta_state
    """
    # For now, we mainly care about the meta_params (update_rule weights)
    # The rnn_state and ema_state are often initialized fresh during
    # agent creation, but we can override them if they're in disco_weights.
    
    updated_state = dict(agent_meta_state)
    
    # If the npz has useful rnn_state or other components, merge them
    # This is a placeholder; adjust based on actual npz structure
    for key in ['rnn_state', 'adv_ema_state', 'td_ema_state']:
        if key in disco_weights:
            updated_state[key] = disco_weights[key]
    
    return updated_state