import numpy as np
import pickle
from pathlib import Path
from .model_loader import get_model

MODEL_DIR = Path(__file__).parent
SCALER_PATH = MODEL_DIR / "scaler.pkl"
ENCODER_PATH = MODEL_DIR / "encoder.pkl"

_scaler = None
_encoder = None

def _load_artifacts():
    """Load scaler and encoder lazily."""
    global _scaler, _encoder
    if _scaler is None:
        if not SCALER_PATH.exists():
             raise FileNotFoundError(f"Scaler file not found at: {SCALER_PATH}")
        with open(SCALER_PATH, "rb") as f:
            _scaler = pickle.load(f)
            
    if _encoder is None:
        if not ENCODER_PATH.exists():
            raise FileNotFoundError(f"Encoder file not found at: {ENCODER_PATH}")
        with open(ENCODER_PATH, "rb") as f:
            _encoder = pickle.load(f)

def _preprocess_features(
    orbital_period_days: float,
    transit_depth_ppm: float,
    transit_duration_hours: float,
    snr: float,
    insolation_flux: float,
) -> np.ndarray:
    """
    Pack features and apply scaling.
    """
    _load_artifacts()
    
    features = np.array([[
        orbital_period_days,
        transit_depth_ppm,
        transit_duration_hours,
        snr,
        insolation_flux
    ]], dtype=np.float32)
    
    return _scaler.transform(features)

def predict(
    orbital_period_days: float,
    transit_depth_ppm: float,
    transit_duration_hours: float,
    snr: float,
    insolation_flux: float,
    return_label: bool = True,
):
    """
    Predict whether an exoplanet is a false positive or candidate.

    Parameters
    ----------
    orbital_period_days : float
    transit_depth_ppm : float
    transit_duration_hours : float
    snr : float
    insolation_flux : float
    return_label : bool, default True
        If True, return a string label ("false_positive"/"candidate").
        If False, return the raw numpy array of probabilities.

    Returns
    -------
    str or np.ndarray
    """
    model = get_model()
    _load_artifacts()
    
    x = _preprocess_features(
        orbital_period_days,
        transit_depth_ppm,
        transit_duration_hours,
        snr,
        insolation_flux
    )
    
    # model.predict output shape is (samples, classes)
    predictions = model.predict(x, verbose=0)[0]
    
    if not return_label:
        return predictions
        
    class_index = np.argmax(predictions)
    # Inverse transform returns an array, take the first item
    return _encoder.inverse_transform([class_index])[0]
