#!/usr/bin/env python3
"""
Simple LightGBM for Aeroclub Flight Ranking 2025 - Memory Optimized
"""

import os
import numpy as np
import pandas as pd
import lightgbm as lgb
from sklearn.preprocessing import LabelEncoder
import warnings
warnings.filterwarnings('ignore')
import gc

print("Starting Simple LightGBM Pipeline")

# Set seed for reproducibility
np.random.seed(42)

# 1. Load data
if os.path.exists('/kaggle/input/aeroclub-recsys-2025/'):
    data_path = '/kaggle/input/aeroclub-recsys-2025/'
else:
    data_path = './'

print("Loading data...")
train_df = pd.read_parquet(f'{data_path}train.parquet')
test_df = pd.read_parquet(f'{data_path}test.parquet')
sample_submission = pd.read_parquet(f'{data_path}sample_submission.parquet')

print(f"Train: {train_df.shape}, Test: {test_df.shape}")

# 2. Simple feature engineering (memory efficient)
def create_simple_features(df):
    df = df.copy()
    
    # Basic price features
    df['price_rank'] = df.groupby('ranker_id')['totalPrice'].rank(method='dense')
    df['price_pct'] = df.groupby('ranker_id')['totalPrice'].rank(pct=True)
    
    # Duration parsing - simplified
    def parse_duration_simple(dur):
        if pd.isna(dur):
            return 300
        try:
            dur_str = str(dur)
            if ':' in dur_str:
                parts = dur_str.replace('.', ':').split(':')
                if len(parts) >= 2:
                    hours = int(parts[-2])
                    minutes = int(parts[-1])
                    return hours * 60 + minutes
            return 300
        except:
            return 300
    
    # Duration features
    if 'legs0_duration' in df.columns:
        df['duration_min'] = df['legs0_duration'].apply(parse_duration_simple)
    else:
        df['duration_min'] = 300
        
    df['duration_rank'] = df.groupby('ranker_id')['duration_min'].rank(method='dense')
    
    # Simple time features
    if 'legs0_departureAt' in df.columns:
        df['legs0_departureAt'] = pd.to_datetime(df['legs0_departureAt'], errors='coerce')
        df['dep_hour'] = df['legs0_departureAt'].dt.hour.fillna(12)
    else:
        df['dep_hour'] = 12
    
    # Route analysis
    df['is_direct'] = (~df['searchRoute'].astype(str).str.contains('/')).astype(int)
    
    # Group size
    df['group_size'] = df.groupby('ranker_id')['Id'].transform('count')
    
    # Simple interaction
    df['price_duration_ratio'] = df['totalPrice'] / (df['duration_min'] + 1)
    
    return df

print("Creating simple features...")
train_feat = create_simple_features(train_df)
del train_df
gc.collect()

test_feat = create_simple_features(test_df)
del test_df
gc.collect()

# 3. Simple feature selection
features = [
    'totalPrice', 'taxes', 'price_rank', 'price_pct',
    'duration_min', 'duration_rank', 'dep_hour', 
    'is_direct', 'group_size', 'price_duration_ratio',
    'pricingInfo_passengerCount'
]

# Ensure all features exist
features = [f for f in features if f in train_feat.columns and f in test_feat.columns]
print(f"Using {len(features)} features")

# Fill missing values
for feat in features:
    train_feat[feat] = train_feat[feat].fillna(0)
    test_feat[feat] = test_feat[feat].fillna(0)

# Prepare data
X_train = train_feat[features]
y_train = train_feat['selected']
groups_train = train_feat['ranker_id']
X_test = test_feat[features]

# Memory optimization - use smaller groups
print("Filtering to manageable groups...")
group_sizes = train_feat.groupby('ranker_id').size()
valid_groups = group_sizes[(group_sizes >= 3) & (group_sizes <= 100)].index  # Much smaller
print(f"Using {len(valid_groups)} valid groups")

mask = train_feat['ranker_id'].isin(valid_groups)
X_train_filt = X_train[mask]
y_train_filt = y_train[mask]
groups_train_filt = groups_train[mask]

# Free memory
del train_feat, X_train, y_train, groups_train
gc.collect()

# 4. Simple model training
print("Training single LightGBM model...")
group_info = groups_train_filt.value_counts(sort=False).sort_index()
query_sizes = group_info.values

params = {
    'objective': 'lambdarank',
    'metric': 'ndcg',
    'ndcg_eval_at': [3],
    'num_leaves': 31,
    'learning_rate': 0.1,
    'feature_fraction': 0.9,
    'bagging_fraction': 0.8,
    'bagging_freq': 5,
    'random_state': 42,
    'verbosity': -1
}

train_data = lgb.Dataset(X_train_filt, label=y_train_filt, group=query_sizes)
model = lgb.train(
    params, 
    train_data, 
    num_boost_round=50,  # Reduced iterations
    valid_sets=[train_data], 
    callbacks=[lgb.log_evaluation(0)]
)

# Free more memory
del X_train_filt, y_train_filt, groups_train_filt, train_data
gc.collect()

# 5. Make predictions
print("Making predictions...")
test_scores = model.predict(X_test)

# 6. Create submission
print("Creating submission...")
submission = test_feat[['Id', 'ranker_id']].copy()
submission['score'] = test_scores

# Rank within each group
submission['selected'] = submission.groupby('ranker_id')['score'].rank(method='first', ascending=False).astype('int64')

# Create final submission
final_submission = submission[['Id', 'ranker_id', 'selected']].copy()
final_submission = final_submission.astype({
    'Id': 'int64',
    'ranker_id': 'object',
    'selected': 'int64'
})

final_submission = final_submission.sort_values('Id').reset_index(drop=True)

# Save submission
final_submission.to_csv('submission.csv', index=False)

print(f"Submission created: {final_submission.shape}")
print("Sample rows:")
print(final_submission.head(10))

print("\nFeature importance:")
feature_imp = pd.DataFrame({
    'feature': features,
    'importance': model.feature_importance(importance_type='gain')
}).sort_values('importance', ascending=False)
print(feature_imp.head(10))

print("\nSubmitting to competition...")
import subprocess
result = subprocess.run([
    'kaggle', 'competitions', 'submit', 
    '-c', 'aeroclub-recsys-2025', 
    '-f', 'submission.csv', 
    '-m', 'Simple LightGBM ranking model'
], capture_output=True, text=True)

print("Submission output:")
print(result.stdout)
if result.stderr:
    print("Submission errors:")
    print(result.stderr)

print("Done!")