#!/usr/bin/env python3
"""
Model classes for crop yield prediction
"""

import pandas as pd
import numpy as np
from typing import Dict, Any, List


class DataPreprocessor:
    """Data preprocessing class that matches the trained model."""
    
    def __init__(self):
        self.label_encoders = {}
        self.scaler = None
        self.imputer = None
        self.feature_names = None
    
    def prepare_features(self, df):
        """Prepare features for prediction."""
        data = df.copy()
        
        # Feature engineering (same as training)
        data['Area_Production_Ratio'] = data['Area'] / (data['Production'] + 1e-6)
        data['Yield_Area_Interaction'] = data.get('Yield', 0) * data['Area']
        data['Production_Per_Area'] = data['Production'] / (data['Area'] + 1e-6)
        
        # Create season dummies
        season_dummies = pd.get_dummies(data['Season'], prefix='Season')
        expected_seasons = ['Season_Autumn', 'Season_Kharif', 'Season_Rabi', 
                          'Season_Summer', 'Season_Total', 'Season_Whole Year', 'Season_Winter']
        for season in expected_seasons:
            if season not in season_dummies.columns:
                season_dummies[season] = 0
        
        data = pd.concat([data, season_dummies[expected_seasons]], axis=1)
        
        # Handle categorical variables
        categorical_cols = ['State', 'District', 'Crop']
        for col in categorical_cols:
            if col in data.columns and col in self.label_encoders:
                unique_values = set(data[col].astype(str))
                known_values = set(self.label_encoders[col].classes_)
                new_values = unique_values - known_values
                
                if new_values:
                    mode_value = self.label_encoders[col].classes_[0]
                    data[col] = data[col].astype(str).replace(list(new_values), mode_value)
                
                data[f'{col}_encoded'] = self.label_encoders[col].transform(data[col].astype(str))
            elif col in data.columns:
                data[f'{col}_encoded'] = pd.Categorical(data[col].astype(str)).codes
        
        # Select features
        feature_cols = ['Crop_Year', 'Area', 'Production', 'Annual_Rainfall', 
                       'Fertilizer', 'Pesticide', 'State_encoded', 'Crop_encoded',
                       'Area_Production_Ratio', 'Yield_Area_Interaction', 
                       'Production_Per_Area'] + expected_seasons
        
        if 'District_encoded' in data.columns:
            feature_cols.append('District_encoded')
        
        available_cols = [col for col in feature_cols if col in data.columns]
        X = data[available_cols].copy()
        
        return X, data
    
    def transform(self, X):
        """Transform features using fitted preprocessors."""
        if self.imputer is None or self.scaler is None:
            raise ValueError("Preprocessor not fitted.")
        
        X_imputed = pd.DataFrame(
            self.imputer.transform(X), 
            columns=X.columns, 
            index=X.index
        )
        
        X_scaled = pd.DataFrame(
            self.scaler.transform(X_imputed), 
            columns=X.columns, 
            index=X.index
        )
        
        return X_scaled