#!/usr/bin/env python3
"""
Dataset Cleaning Script for Combined Crop Data
This script cleans and improves the quality of the combined_crop_data.csv dataset
by handling missing values, zero values, and ensuring data consistency.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

def load_and_analyze_data(file_path):
    """Load the dataset and perform initial analysis"""
    print("Loading dataset...")
    df = pd.read_csv(file_path)
    
    print(f"Dataset shape: {df.shape}")
    print(f"Columns: {list(df.columns)}")
    print("\nFirst few rows:")
    print(df.head())
    
    print("\nDataset info:")
    print(df.info())
    
    print("\nMissing values count:")
    missing_counts = df.isnull().sum()
    print(missing_counts[missing_counts > 0])
    
    print("\nMissing values percentage:")
    missing_percentages = (df.isnull().sum() / len(df)) * 100
    print(missing_percentages[missing_percentages > 0])
    
    return df

def analyze_zero_values(df):
    """Analyze zero values in numerical columns"""
    print("\n" + "="*50)
    print("ANALYZING ZERO VALUES")
    print("="*50)
    
    numerical_columns = ['Area', 'Production', 'Annual_Rainfall', 'Fertilizer', 'Pesticide', 'Yield']
    
    for col in numerical_columns:
        if col in df.columns:
            zero_count = (df[col] == 0).sum()
            zero_percentage = (zero_count / len(df)) * 100
            print(f"{col}: {zero_count} zeros ({zero_percentage:.2f}%)")
            
            if zero_count > 0:
                print(f"  Sample rows with zero {col}:")
                sample_zeros = df[df[col] == 0][['Crop', 'State', 'District', col]].head(3)
                print(f"  {sample_zeros.to_string()}")
                print()

def clean_missing_values(df):
    """Handle missing values intelligently"""
    print("\n" + "="*50)
    print("CLEANING MISSING VALUES")
    print("="*50)
    
    df_cleaned = df.copy()
    
    # Handle missing Annual_Rainfall
    if 'Annual_Rainfall' in df_cleaned.columns:
        missing_rainfall = df_cleaned['Annual_Rainfall'].isnull().sum()
        print(f"Missing Annual_Rainfall values: {missing_rainfall}")
        
        if missing_rainfall > 0:
            # Fill with median rainfall by state and season
            df_cleaned['Annual_Rainfall'] = df_cleaned.groupby(['State', 'Season'])['Annual_Rainfall'].transform(
                lambda x: x.fillna(x.median())
            )
            
            # If still missing, fill with overall median
            overall_median = df_cleaned['Annual_Rainfall'].median()
            df_cleaned['Annual_Rainfall'].fillna(overall_median, inplace=True)
            print(f"Filled missing Annual_Rainfall with median values")
    
    # Handle missing Fertilizer
    if 'Fertilizer' in df_cleaned.columns:
        missing_fertilizer = df_cleaned['Fertilizer'].isnull().sum()
        print(f"Missing Fertilizer values: {missing_fertilizer}")
        
        if missing_fertilizer > 0:
            # Fill with median fertilizer by crop and state
            df_cleaned['Fertilizer'] = df_cleaned.groupby(['Crop', 'State'])['Fertilizer'].transform(
                lambda x: x.fillna(x.median())
            )
            
            # If still missing, fill with crop median
            df_cleaned['Fertilizer'] = df_cleaned.groupby('Crop')['Fertilizer'].transform(
                lambda x: x.fillna(x.median())
            )
            
            # If still missing, fill with overall median
            overall_median = df_cleaned['Fertilizer'].median()
            df_cleaned['Fertilizer'].fillna(overall_median, inplace=True)
            print(f"Filled missing Fertilizer with median values")
    
    # Handle missing Pesticide
    if 'Pesticide' in df_cleaned.columns:
        missing_pesticide = df_cleaned['Pesticide'].isnull().sum()
        print(f"Missing Pesticide values: {missing_pesticide}")
        
        if missing_pesticide > 0:
            # Fill with median pesticide by crop and state
            df_cleaned['Pesticide'] = df_cleaned.groupby(['Crop', 'State'])['Pesticide'].transform(
                lambda x: x.fillna(x.median())
            )
            
            # If still missing, fill with crop median
            df_cleaned['Pesticide'] = df_cleaned.groupby('Crop')['Pesticide'].transform(
                lambda x: x.fillna(x.median())
            )
            
            # If still missing, fill with overall median
            overall_median = df_cleaned['Pesticide'].median()
            df_cleaned['Pesticide'].fillna(overall_median, inplace=True)
            print(f"Filled missing Pesticide with median values")
    
    return df_cleaned

def handle_zero_values(df):
    """Handle zero values appropriately"""
    print("\n" + "="*50)
    print("HANDLING ZERO VALUES")
    print("="*50)
    
    df_cleaned = df.copy()
    
    # For Area and Production, zero might be legitimate (no cultivation)
    # But we need to be careful about yield calculation
    
    # Handle problematic zero Area values where Production > 0
    problematic_area = (df_cleaned['Area'] == 0) & (df_cleaned['Production'] > 0)
    if problematic_area.sum() > 0:
        print(f"Found {problematic_area.sum()} records with Area=0 but Production>0")
        
        # Calculate median area per unit production for each crop
        area_prod_ratio = df_cleaned[df_cleaned['Area'] > 0].groupby('Crop').apply(
            lambda x: (x['Area'] / x['Production']).median()
        ).to_dict()
        
        for idx in df_cleaned[problematic_area].index:
            crop = df_cleaned.loc[idx, 'Crop']
            production = df_cleaned.loc[idx, 'Production']
            if crop in area_prod_ratio:
                estimated_area = production * area_prod_ratio[crop]
                df_cleaned.loc[idx, 'Area'] = estimated_area
                print(f"  Fixed Area for {crop}: estimated {estimated_area:.3f} based on production")
    
    # Handle zero Production where Area > 0 (crop failure cases)
    zero_production = (df_cleaned['Production'] == 0) & (df_cleaned['Area'] > 0)
    if zero_production.sum() > 0:
        print(f"Found {zero_production.sum()} records with Production=0 but Area>0 (possible crop failures)")
        # These might be legitimate (crop failures), so we'll keep them but ensure yield is 0
        df_cleaned.loc[zero_production, 'Yield'] = 0
    
    return df_cleaned

def recalculate_yield(df):
    """Recalculate yield to ensure consistency"""
    print("\n" + "="*50)
    print("RECALCULATING YIELD FOR CONSISTENCY")
    print("="*50)
    
    df_cleaned = df.copy()
    
    # Calculate yield where both area and production are non-zero
    mask = (df_cleaned['Area'] > 0) & (df_cleaned['Production'] > 0)
    
    # Store original yield for comparison
    original_yield = df_cleaned['Yield'].copy()
    
    # Recalculate yield as Production/Area (assuming Production is in appropriate units)
    df_cleaned.loc[mask, 'Yield'] = df_cleaned.loc[mask, 'Production'] / df_cleaned.loc[mask, 'Area']
    
    # For records where area or production is zero, set yield to 0
    zero_mask = (df_cleaned['Area'] == 0) | (df_cleaned['Production'] == 0)
    df_cleaned.loc[zero_mask, 'Yield'] = 0
    
    # Compare with original yields
    yield_diff = abs(df_cleaned['Yield'] - original_yield)
    significant_changes = yield_diff > 0.01  # More than 1% difference
    
    print(f"Yield recalculated for {mask.sum()} records")
    print(f"Significant changes in yield: {significant_changes.sum()} records")
    
    return df_cleaned

def add_explicit_units(df):
    """Add explicit units to yield and production fields"""
    print("\n" + "="*50)
    print("ADDING EXPLICIT UNITS")
    print("="*50)
    
    df_cleaned = df.copy()
    
    # Rename columns to include units
    column_renames = {}
    
    if 'Yield' in df_cleaned.columns:
        column_renames['Yield'] = 'Yield_kg_per_hectare'
        print("Renamed 'Yield' to 'Yield_kg_per_hectare'")
    
    if 'Production' in df_cleaned.columns:
        # Assuming production is in tons (common for agricultural data)
        column_renames['Production'] = 'Production_tons'
        print("Renamed 'Production' to 'Production_tons'")
    
    if 'Area' in df_cleaned.columns:
        column_renames['Area'] = 'Area_hectares'
        print("Renamed 'Area' to 'Area_hectares'")
    
    if 'Annual_Rainfall' in df_cleaned.columns:
        column_renames['Annual_Rainfall'] = 'Annual_Rainfall_mm'
        print("Renamed 'Annual_Rainfall' to 'Annual_Rainfall_mm'")
    
    if 'Fertilizer' in df_cleaned.columns:
        column_renames['Fertilizer'] = 'Fertilizer_kg_per_hectare'
        print("Renamed 'Fertilizer' to 'Fertilizer_kg_per_hectare'")
    
    if 'Pesticide' in df_cleaned.columns:
        column_renames['Pesticide'] = 'Pesticide_kg_per_hectare'
        print("Renamed 'Pesticide' to 'Pesticide_kg_per_hectare'")
    
    df_cleaned.rename(columns=column_renames, inplace=True)
    
    return df_cleaned

def detect_and_handle_outliers(df):
    """Detect and handle outliers in numerical columns"""
    print("\n" + "="*50)
    print("DETECTING AND HANDLING OUTLIERS")
    print("="*50)
    
    df_cleaned = df.copy()
    numerical_cols = df_cleaned.select_dtypes(include=[np.number]).columns
    
    outliers_summary = {}
    
    for col in numerical_cols:
        if col not in ['Crop_Year']:  # Skip year column
            Q1 = df_cleaned[col].quantile(0.25)
            Q3 = df_cleaned[col].quantile(0.75)
            IQR = Q3 - Q1
            lower_bound = Q1 - 1.5 * IQR
            upper_bound = Q3 + 1.5 * IQR
            
            outliers = (df_cleaned[col] < lower_bound) | (df_cleaned[col] > upper_bound)
            outliers_count = outliers.sum()
            
            if outliers_count > 0:
                outliers_summary[col] = {
                    'count': outliers_count,
                    'percentage': (outliers_count / len(df_cleaned)) * 100,
                    'lower_bound': lower_bound,
                    'upper_bound': upper_bound
                }
                
                # For extreme outliers, cap them at reasonable bounds
                extreme_outliers = (df_cleaned[col] > upper_bound + 2 * IQR) | (df_cleaned[col] < lower_bound - 2 * IQR)
                if extreme_outliers.sum() > 0:
                    print(f"Capping {extreme_outliers.sum()} extreme outliers in {col}")
                    df_cleaned.loc[df_cleaned[col] > upper_bound + 2 * IQR, col] = upper_bound
                    df_cleaned.loc[df_cleaned[col] < lower_bound - 2 * IQR, col] = max(0, lower_bound)
    
    print("\nOutliers summary:")
    for col, info in outliers_summary.items():
        print(f"{col}: {info['count']} outliers ({info['percentage']:.2f}%)")
    
    return df_cleaned

def ensure_data_consistency(df):
    """Ensure logical consistency between related fields"""
    print("\n" + "="*50)
    print("ENSURING DATA CONSISTENCY")
    print("="*50)
    
    df_cleaned = df.copy()
    
    # Ensure consistent data types
    if 'Crop_Year' in df_cleaned.columns:
        df_cleaned['Crop_Year'] = df_cleaned['Crop_Year'].astype(int)
    
    # Ensure non-negative values for physical quantities
    numerical_cols = df_cleaned.select_dtypes(include=[np.number]).columns
    for col in numerical_cols:
        if col != 'Crop_Year':
            negative_count = (df_cleaned[col] < 0).sum()
            if negative_count > 0:
                print(f"Found {negative_count} negative values in {col}, setting to 0")
                df_cleaned.loc[df_cleaned[col] < 0, col] = 0
    
    # Standardize text fields
    text_columns = ['Crop', 'Season', 'State', 'District']
    for col in text_columns:
        if col in df_cleaned.columns:
            df_cleaned[col] = df_cleaned[col].str.strip()  # Remove whitespace
            df_cleaned[col] = df_cleaned[col].str.title()  # Standardize capitalization
    
    return df_cleaned

def generate_data_quality_report(original_df, cleaned_df):
    """Generate a comprehensive data quality report"""
    print("\n" + "="*60)
    print("DATA QUALITY IMPROVEMENT REPORT")
    print("="*60)
    
    print(f"\nDataset size: {len(cleaned_df)} records")
    
    # Missing values comparison
    print("\nMissing Values - Before vs After:")
    print("-" * 40)
    for col in original_df.columns:
        original_missing = original_df[col].isnull().sum()
        cleaned_missing = cleaned_df[col].isnull().sum() if col in cleaned_df.columns else 0
        if original_missing > 0 or cleaned_missing > 0:
            print(f"{col:20} {original_missing:8} → {cleaned_missing:8}")
    
    # Zero values comparison (for numerical columns)
    print("\nZero Values - Before vs After:")
    print("-" * 40)
    numerical_cols = ['Area', 'Production', 'Annual_Rainfall', 'Fertilizer', 'Pesticide', 'Yield']
    for col in numerical_cols:
        if col in original_df.columns:
            original_zeros = (original_df[col] == 0).sum()
            
            # Find corresponding column in cleaned dataset
            cleaned_col = None
            for cleaned_column in cleaned_df.columns:
                if col.lower() in cleaned_column.lower():
                    cleaned_col = cleaned_column
                    break
            
            if cleaned_col:
                cleaned_zeros = (cleaned_df[cleaned_col] == 0).sum()
                print(f"{col:20} {original_zeros:8} → {cleaned_zeros:8}")
    
    # Data quality metrics
    print("\nData Quality Improvements:")
    print("-" * 40)
    
    # Completeness
    original_completeness = (1 - original_df.isnull().sum().sum() / (len(original_df) * len(original_df.columns))) * 100
    cleaned_completeness = (1 - cleaned_df.isnull().sum().sum() / (len(cleaned_df) * len(cleaned_df.columns))) * 100
    print(f"Completeness:       {original_completeness:.2f}% → {cleaned_completeness:.2f}%")
    
    # Column names with units
    print("\nColumns with explicit units:")
    print("-" * 40)
    for col in cleaned_df.columns:
        if any(unit in col.lower() for unit in ['kg', 'hectare', 'tons', 'mm']):
            print(f"✓ {col}")
    
    return cleaned_df

def save_cleaned_dataset(df, output_path):
    """Save the cleaned dataset"""
    print(f"\nSaving cleaned dataset to: {output_path}")
    df.to_csv(output_path, index=False)
    print(f"✓ Cleaned dataset saved with {len(df)} records and {len(df.columns)} columns")
    
    # Also save a sample for verification
    sample_path = output_path.replace('.csv', '_sample.csv')
    df.head(1000).to_csv(sample_path, index=False)
    print(f"✓ Sample dataset saved to: {sample_path}")

def main():
    """Main function to orchestrate the data cleaning process"""
    print("="*60)
    print("CROP DATA CLEANING AND QUALITY IMPROVEMENT")
    print("="*60)
    
    # File paths
    input_file = Path("/home/aiavid/Yeild_pred_SIH/data/combined_crop_data.csv")
    output_file = Path("/home/aiavid/Yeild_pred_SIH/data/combined_crop_data_cleaned.csv")
    
    # Load and analyze original data
    original_df = load_and_analyze_data(input_file)
    
    # Analyze issues
    analyze_zero_values(original_df)
    
    # Start cleaning process
    print("\n" + "="*60)
    print("STARTING DATA CLEANING PROCESS")
    print("="*60)
    
    # Step 1: Handle missing values
    cleaned_df = clean_missing_values(original_df)
    
    # Step 2: Handle zero values
    cleaned_df = handle_zero_values(cleaned_df)
    
    # Step 3: Recalculate yield for consistency
    cleaned_df = recalculate_yield(cleaned_df)
    
    # Step 4: Add explicit units
    cleaned_df = add_explicit_units(cleaned_df)
    
    # Step 5: Handle outliers
    cleaned_df = detect_and_handle_outliers(cleaned_df)
    
    # Step 6: Ensure data consistency
    cleaned_df = ensure_data_consistency(cleaned_df)
    
    # Generate quality report
    cleaned_df = generate_data_quality_report(original_df, cleaned_df)
    
    # Save cleaned dataset
    save_cleaned_dataset(cleaned_df, output_file)
    
    print("\n" + "="*60)
    print("DATA CLEANING COMPLETED SUCCESSFULLY!")
    print("="*60)
    
    return cleaned_df

if __name__ == "__main__":
    main()