"""
Feature Extraction Module for Land Classification
Chuẩn hóa việc trích xuất features từ satellite data cho cả training và prediction
"""

import numpy as np
import xarray as xr
from typing import List, Dict, Tuple, Optional


class FeatureExtractor:
    """
    Extract features từ Sentinel-2 và Sentinel-1 data
    Hỗ trợ 2 modes:
    - 'simple': 3 features cơ bản (NDVI_mean, VH_mean, VV_mean)
    - 'temporal': 39 features time-series (NDVI + NDWI + NDBI theo thời gian)
    """
    
    FEATURE_MODES = {
        'simple': {
            'n_features': 3,
            'features': ['NDVI_mean', 'VH_db_mean', 'VV_db_mean'],
            'description': 'Simple aggregate features (mean only)'
        },
        'temporal': {
            'n_features': 39,
            'features': None,  # Generated dynamically based on time steps
            'description': 'Temporal features with NDVI, NDWI, NDBI time series'
        },
        'extended': {
            'n_features': 15,
            'features': [
                'NDVI_mean', 'NDVI_std', 'NDVI_min', 'NDVI_max',
                'NDWI_mean', 'NDWI_std', 'NDWI_min', 'NDWI_max',
                'NDBI_mean', 'NDBI_std', 'NDBI_min', 'NDBI_max',
                'VH_db_mean', 'VV_db_mean', 'VH_VV_ratio'
            ],
            'description': 'Extended aggregate features with statistics'
        },
        'odc': {
            'n_features': 8,
            'features': [
                'ndvi_mean', 'ndvi_min', 'ndvi_max', 'ndvi_std', 'ndvi_range',
                'ndwi_mean', 'ndbi_mean', 'evi_mean'
            ],
            'description': 'ODC mode: 8 aggregate features (NDVI stats + NDWI/NDBI/EVI mean) - matches 01.train_ODC.ipynb'
        }
    }
    
    def __init__(self, mode: str = 'simple'):
        """
        Initialize FeatureExtractor
        
        Args:
            mode: 'simple', 'temporal', hoặc 'extended'
        """
        if mode not in self.FEATURE_MODES:
            raise ValueError(f"Invalid mode: {mode}. Choose from {list(self.FEATURE_MODES.keys())}")
        
        self.mode = mode
        self.config = self.FEATURE_MODES[mode]
    
    def get_feature_names(self, n_timesteps: Optional[int] = None) -> List[str]:
        """
        Lấy danh sách tên features
        
        Args:
            n_timesteps: Số timesteps (chỉ cần cho mode='temporal')
        
        Returns:
            List tên features
        """
        if self.mode == 'temporal':
            if n_timesteps is None:
                raise ValueError("n_timesteps required for temporal mode")
            
            features = []
            # NDVI time series
            for t in range(n_timesteps):
                features.append(f'NDVI_t{t+1}')
            # NDWI time series
            for t in range(n_timesteps):
                features.append(f'NDWI_t{t+1}')
            # NDBI time series
            for t in range(n_timesteps):
                features.append(f'NDBI_t{t+1}')
            
            # VH/VV radar (mean across time)
            features.append('VH_db_mean')
            features.append('VV_db_mean')
            features.append('VH_VV_ratio')
            
            return features
        else:
            return self.config['features']
    
    def extract_simple_features(
        self, 
        ndvi_data: xr.DataArray,
        vh_data: Optional[xr.DataArray] = None,
        vv_data: Optional[xr.DataArray] = None
    ) -> np.ndarray:
        """
        Extract simple features (3 features: NDVI_mean, VH_db_mean, VV_db_mean)
        
        Args:
            ndvi_data: NDVI DataArray (có thể có time dimension)
            vh_data: VH radar DataArray
            vv_data: VV radar DataArray
        
        Returns:
            Feature array shape (n_pixels, 3)
        """
        # Calculate NDVI mean
        if 'time' in ndvi_data.dims:
            ndvi_mean = ndvi_data.mean(dim='time')
        else:
            ndvi_mean = ndvi_data
        
        # Flatten to pixels
        ndvi_flat = ndvi_mean.values.flatten()
        
        # Calculate radar features if available
        if vh_data is not None and vv_data is not None:
            if 'time' in vh_data.dims:
                vh_mean = vh_data.mean(dim='time')
                vv_mean = vv_data.mean(dim='time')
            else:
                vh_mean = vh_data
                vv_mean = vv_data
            
            vh_flat = vh_mean.values.flatten()
            vv_flat = vv_mean.values.flatten()
        else:
            # If no radar data, use zeros
            vh_flat = np.zeros_like(ndvi_flat)
            vv_flat = np.zeros_like(ndvi_flat)
        
        # Stack features
        features = np.column_stack([ndvi_flat, vh_flat, vv_flat])
        
        return features
    
    def extract_temporal_features(
        self,
        s2_data: xr.Dataset,
        vh_data: Optional[xr.DataArray] = None,
        vv_data: Optional[xr.DataArray] = None
    ) -> np.ndarray:
        """
        Extract temporal features (39 features: time series của NDVI, NDWI, NDBI + radar)
        
        Args:
            s2_data: Sentinel-2 Dataset với bands B02, B03, B04, B08, B11
            vh_data: VH radar DataArray
            vv_data: VV radar DataArray
        
        Returns:
            Feature array shape (n_pixels, 39)
        """
        # Calculate spectral indices
        nir = s2_data["B08"].astype('float32')
        red = s2_data["B04"].astype('float32')
        green = s2_data["B03"].astype('float32')
        swir = s2_data["B11"].astype('float32') if "B11" in s2_data else s2_data["B02"]  # Fallback to B02
        
        # NDVI = (NIR - Red) / (NIR + Red)
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # NDWI = (Green - NIR) / (Green + NIR)
        ndwi = (green - nir) / (green + nir + 1e-8)
        
        # NDBI = (SWIR - NIR) / (SWIR + NIR)
        ndbi = (swir - nir) / (swir + nir + 1e-8)
        
        # Resample to monthly if time dimension exists
        if 'time' in ndvi.dims:
            ndvi_monthly = ndvi.resample(time="1ME").mean()
            ndwi_monthly = ndwi.resample(time="1ME").mean()
            ndbi_monthly = ndbi.resample(time="1ME").mean()
        else:
            ndvi_monthly = ndvi
            ndwi_monthly = ndwi
            ndbi_monthly = ndbi
        
        # Get dimensions
        n_times = len(ndvi_monthly.time) if 'time' in ndvi_monthly.dims else 1
        y_size = len(ndvi_monthly.y)
        x_size = len(ndvi_monthly.x)
        n_pixels = y_size * x_size
        
        # Extract temporal features
        features_list = []
        
        # NDVI time series
        for t in range(n_times):
            if 'time' in ndvi_monthly.dims:
                ndvi_t = ndvi_monthly.isel(time=t).values.flatten()
            else:
                ndvi_t = ndvi_monthly.values.flatten()
            features_list.append(ndvi_t)
        
        # NDWI time series
        for t in range(n_times):
            if 'time' in ndwi_monthly.dims:
                ndwi_t = ndwi_monthly.isel(time=t).values.flatten()
            else:
                ndwi_t = ndwi_monthly.values.flatten()
            features_list.append(ndwi_t)
        
        # NDBI time series
        for t in range(n_times):
            if 'time' in ndbi_monthly.dims:
                ndbi_t = ndbi_monthly.isel(time=t).values.flatten()
            else:
                ndbi_t = ndbi_monthly.values.flatten()
            features_list.append(ndbi_t)
        
        # Stack all spectral features
        features = np.column_stack(features_list)
        
        # Add radar features if available
        if vh_data is not None and vv_data is not None:
            if 'time' in vh_data.dims:
                vh_mean = vh_data.mean(dim='time')
                vv_mean = vv_data.mean(dim='time')
            else:
                vh_mean = vh_data
                vv_mean = vv_data
            
            vh_flat = vh_mean.values.flatten()
            vv_flat = vv_mean.values.flatten()
            vh_vv_ratio = vh_flat / (vv_flat + 1e-8)
            
            # Add radar features
            features = np.column_stack([features, vh_flat, vv_flat, vh_vv_ratio])
        
        return features
    
    def extract_odc_features(
        self,
        s2_data: xr.Dataset,
        vh_data: Optional[xr.DataArray] = None,
        vv_data: Optional[xr.DataArray] = None
    ) -> np.ndarray:
        """
        Extract ODC aggregate features (8 features matching 01.train_ODC.ipynb):
        ndvi_mean, ndvi_min, ndvi_max, ndvi_std, ndvi_range, ndwi_mean, ndbi_mean, evi_mean
        
        Args:
            s2_data: Sentinel-2 Dataset with B02, B03, B04, B08, B11
            vh_data: Not used in ODC mode
            vv_data: Not used in ODC mode
        
        Returns:
            Feature array shape (n_pixels, 8)
        """
        # Calculate spectral indices
        nir = s2_data["B08"].astype('float32')
        red = s2_data["B04"].astype('float32')
        green = s2_data["B03"].astype('float32')
        blue = s2_data["B02"].astype('float32')
        swir = s2_data["B11"].astype('float32') if "B11" in s2_data else s2_data["B02"]
        
        # NDVI = (NIR - Red) / (NIR + Red)
        ndvi = (nir - red) / (nir + red + 1e-8)
        
        # NDWI = (Green - NIR) / (Green + NIR)
        ndwi = (green - nir) / (green + nir + 1e-8)
        
        # NDBI = (SWIR - NIR) / (SWIR + NIR)
        ndbi = (swir - nir) / (swir + nir + 1e-8)
        
        # EVI = 2.5 * (NIR - Red) / (NIR + 6*Red - 7.5*Blue + 1)
        evi = 2.5 * (nir - red) / (nir + 6*red - 7.5*blue + 1)
        
        features_list = []
        
        # NDVI statistics (5 features)
        if 'time' in ndvi.dims:
            features_list.append(ndvi.mean(dim='time').values.flatten())  # ndvi_mean
            features_list.append(ndvi.min(dim='time').values.flatten())   # ndvi_min
            features_list.append(ndvi.max(dim='time').values.flatten())   # ndvi_max
            features_list.append(ndvi.std(dim='time').values.flatten())   # ndvi_std
            ndvi_range = (ndvi.max(dim='time') - ndvi.min(dim='time')).values.flatten()
            features_list.append(ndvi_range)  # ndvi_range
        else:
            ndvi_flat = ndvi.values.flatten()
            features_list.extend([ndvi_flat, ndvi_flat, ndvi_flat, np.zeros_like(ndvi_flat), np.zeros_like(ndvi_flat)])
        
        # NDWI mean (1 feature)
        if 'time' in ndwi.dims:
            features_list.append(ndwi.mean(dim='time').values.flatten())  # ndwi_mean
        else:
            features_list.append(ndwi.values.flatten())
        
        # NDBI mean (1 feature)
        if 'time' in ndbi.dims:
            features_list.append(ndbi.mean(dim='time').values.flatten())  # ndbi_mean
        else:
            features_list.append(ndbi.values.flatten())
        
        # EVI mean (1 feature)
        if 'time' in evi.dims:
            features_list.append(evi.mean(dim='time').values.flatten())  # evi_mean
        else:
            features_list.append(evi.values.flatten())
        
        # Stack all features (total: 8 features)
        features = np.column_stack(features_list)
        
        return features
    
    def extract_extended_features(
        self,
        s2_data: xr.Dataset,
        vh_data: Optional[xr.DataArray] = None,
        vv_data: Optional[xr.DataArray] = None
    ) -> np.ndarray:
        """
        Extract extended aggregate features (15 features: stats của NDVI, NDWI, NDBI + radar)
        
        Args:
            s2_data: Sentinel-2 Dataset
            vh_data: VH radar DataArray
            vv_data: VV radar DataArray
        
        Returns:
            Feature array shape (n_pixels, 15)
        """
        # Calculate spectral indices
        nir = s2_data["B08"].astype('float32')
        red = s2_data["B04"].astype('float32')
        green = s2_data["B03"].astype('float32')
        swir = s2_data["B11"].astype('float32') if "B11" in s2_data else s2_data["B02"]
        
        ndvi = (nir - red) / (nir + red + 1e-8)
        ndwi = (green - nir) / (green + nir + 1e-8)
        ndbi = (swir - nir) / (swir + nir + 1e-8)
        
        features_list = []
        
        # NDVI statistics
        if 'time' in ndvi.dims:
            features_list.append(ndvi.mean(dim='time').values.flatten())
            features_list.append(ndvi.std(dim='time').values.flatten())
            features_list.append(ndvi.min(dim='time').values.flatten())
            features_list.append(ndvi.max(dim='time').values.flatten())
        else:
            ndvi_flat = ndvi.values.flatten()
            features_list.extend([ndvi_flat, np.zeros_like(ndvi_flat), ndvi_flat, ndvi_flat])
        
        # NDWI statistics
        if 'time' in ndwi.dims:
            features_list.append(ndwi.mean(dim='time').values.flatten())
            features_list.append(ndwi.std(dim='time').values.flatten())
            features_list.append(ndwi.min(dim='time').values.flatten())
            features_list.append(ndwi.max(dim='time').values.flatten())
        else:
            ndwi_flat = ndwi.values.flatten()
            features_list.extend([ndwi_flat, np.zeros_like(ndwi_flat), ndwi_flat, ndwi_flat])
        
        # NDBI statistics
        if 'time' in ndbi.dims:
            features_list.append(ndbi.mean(dim='time').values.flatten())
            features_list.append(ndbi.std(dim='time').values.flatten())
            features_list.append(ndbi.min(dim='time').values.flatten())
            features_list.append(ndbi.max(dim='time').values.flatten())
        else:
            ndbi_flat = ndbi.values.flatten()
            features_list.extend([ndbi_flat, np.zeros_like(ndbi_flat), ndbi_flat, ndbi_flat])
        
        # Stack spectral features
        features = np.column_stack(features_list)
        
        # Add radar features
        if vh_data is not None and vv_data is not None:
            if 'time' in vh_data.dims:
                vh_mean = vh_data.mean(dim='time')
                vv_mean = vv_data.mean(dim='time')
            else:
                vh_mean = vh_data
                vv_mean = vv_data
            
            vh_flat = vh_mean.values.flatten()
            vv_flat = vv_mean.values.flatten()
            vh_vv_ratio = vh_flat / (vv_flat + 1e-8)
            
            features = np.column_stack([features, vh_flat, vv_flat, vh_vv_ratio])
        
        return features
    
    def extract(
        self,
        s2_data: Optional[xr.Dataset] = None,
        ndvi_data: Optional[xr.DataArray] = None,
        vh_data: Optional[xr.DataArray] = None,
        vv_data: Optional[xr.DataArray] = None
    ) -> np.ndarray:
        """
        Extract features theo mode đã chọn
        
        Args:
            s2_data: Sentinel-2 Dataset (cần cho temporal, extended, và odc modes)
            ndvi_data: NDVI DataArray (cần cho simple mode)
            vh_data: VH radar DataArray
            vv_data: VV radar DataArray
        
        Returns:
            Feature array
        """
        if self.mode == 'simple':
            if ndvi_data is None:
                raise ValueError("ndvi_data required for simple mode")
            return self.extract_simple_features(ndvi_data, vh_data, vv_data)
        
        elif self.mode == 'temporal':
            if s2_data is None:
                raise ValueError("s2_data required for temporal mode")
            return self.extract_temporal_features(s2_data, vh_data, vv_data)
        
        elif self.mode == 'extended':
            if s2_data is None:
                raise ValueError("s2_data required for extended mode")
            return self.extract_extended_features(s2_data, vh_data, vv_data)
        
        elif self.mode == 'odc':
            if s2_data is None:
                raise ValueError("s2_data required for odc mode")
            return self.extract_odc_features(s2_data, vh_data, vv_data)
        
        else:
            raise ValueError(f"Unknown mode: {self.mode}")
    
    def get_info(self) -> Dict:
        """Lấy thông tin về feature extraction mode"""
        return {
            'mode': self.mode,
            'n_features': self.config['n_features'],
            'description': self.config['description']
        }


def get_feature_extractor(mode: str = 'simple') -> FeatureExtractor:
    """
    Factory function để tạo FeatureExtractor
    
    Args:
        mode: 'simple', 'temporal', 'extended', hoặc 'odc'
    
    Returns:
        FeatureExtractor instance
    """
    return FeatureExtractor(mode=mode)
