from kfp import dsl
from kfp import compiler

# -------------------------
# COMPONENT 1: Download + Prepare Taxi Dataset
# -------------------------
@dsl.component(
    base_image='python:3.9',
    packages_to_install=['pandas', 'kaggle']
)
def prepare_taxi_data(output_csv: dsl.Output[dsl.Dataset]):
    import pandas as pd
    import os
    import random
    from zipfile import ZipFile

    # Ensure Kaggle API credentials are configured
    os.environ['KAGGLE_USERNAME'] = 'bhargavisatya27'
    os.environ['KAGGLE_KEY'] = '54897dd35575cfdeb71baf9a17c85b17'

    # Download dataset
    os.system("kaggle datasets download -d adelanseur/taxi-trips-chicago-2024 -p .")

    # Extract zip
    for file in os.listdir('.'):
        if file.endswith('.zip'):
            with ZipFile(file, 'r') as zip_ref:
                zip_ref.extractall('.')

    # Load the main data file (find .csv file automatically)
    csv_file = [f for f in os.listdir('.') if f.endswith('.csv')][0]
    df = pd.read_csv(csv_file)

    # Select relevant columns (adapt based on Kaggle column names)
    # Usually columns like: 'Pickup Community Area', 'Dropoff Community Area', 'Trip Total'
    # You may need to adjust these if the Kaggle file uses different names
    selected_cols = ['Pickup Community Area', 'Dropoff Community Area', 'Trip Total']
    df_small = df[selected_cols].rename(columns={
        'Pickup Community Area': 'start_location',
        'Dropoff Community Area': 'end_location',
        'Trip Total': 'price'
    })

    # Drop missing values and sample 2000 rows
    df_small = df_small.dropna().sample(5000, random_state=42)

    # Save to output dataset
    df_small.to_csv(output_csv.path, index=False)
    print(f"✅ Prepared taxi dataset saved to: {output_csv.path}")

# -------------------------
# COMPONENT 2: Train Fare Model
# -------------------------
@dsl.component(
    base_image='python:3.9',
    packages_to_install=['pandas', 'scikit-learn', 'joblib']
)
def train_fare_model(dataset: dsl.Input[dsl.Dataset], model: dsl.Output[dsl.Model]):
    import pandas as pd
    import joblib
    from sklearn.preprocessing import OneHotEncoder
    from sklearn.ensemble import RandomForestRegressor
    from sklearn.compose import ColumnTransformer
    from sklearn.pipeline import Pipeline

    df = pd.read_csv(dataset.path)
    X = df[['start_location', 'end_location']]
    y = df['price']

    # One-hot encode start and end locations
    categorical_features = ['start_location', 'end_location']
    preprocessor = ColumnTransformer(
        transformers=[('cat', OneHotEncoder(handle_unknown='ignore'), categorical_features)]
    )

    # Random Forest Regressor pipeline
    model_pipeline = Pipeline(steps=[
        ('preprocessor', preprocessor),
        ('regressor', RandomForestRegressor(n_estimators=100, random_state=42))
    ])

    model_pipeline.fit(X, y)
    joblib.dump(model_pipeline, model.path)
    print(f"✅ Model trained and saved to: {model.path}")

# -------------------------
# COMPONENT 3: Evaluate Model
# -------------------------
@dsl.component(
    base_image='python:3.9',
    packages_to_install=['pandas', 'scikit-learn', 'joblib']
)
def evaluate_model(dataset: dsl.Input[dsl.Dataset], model: dsl.Input[dsl.Model]) -> float:
    import pandas as pd
    import joblib
    from sklearn.metrics import r2_score

    df = pd.read_csv(dataset.path)
    X = df[['start_location', 'end_location']]
    y_true = df['price']

    model_pipeline = joblib.load(model.path)
    y_pred = model_pipeline.predict(X)
    r2 = r2_score(y_true, y_pred)

    print(f"✅ Model R² score: {r2}")
    return float(r2)

# -------------------------
# PIPELINE DEFINITION
# -------------------------
@dsl.pipeline(
    name='chicago-taxi-fare-pipeline',
    description='Pipeline using Kaggle Chicago Taxi Trips 2024 to predict fare from start/end locations.'
)
def taxi_pipeline():
    data = prepare_taxi_data()
    model = train_fare_model(dataset=data.outputs['output_csv'])
    evaluate_model(dataset=data.outputs['output_csv'], model=model.outputs['model'])

# -------------------------
# COMPILE THE PIPELINE
# -------------------------
if __name__ == '__main__':
    compiler.Compiler().compile(
        pipeline_func=taxi_pipeline,
        package_path='chicago_taxi_pipeline.yaml'
    )
    print("✅ Pipeline compiled successfully: chicago_taxi_pipeline.yaml")

    
  

    
      # Check pipeline runs:
kubectl get pipelineruns -n kubeflow

# View specific run details:
kubectl describe pipelinerun <run-name> -n kubeflow

# Check pod logs:
kubectl logs -f <pod-name> -n kubeflow