# --- Imports ---
import pandas as pd
from econml.dml import CausalForestDML
import numpy as np

# --- Load your data ---
df = pd.read_excel("Thereshold.xlsx")
df.columns = df.columns.str.strip()  # remove whitespace

# --- Encode treatments ---
df['Treatments_encoded'] = df['Treatments'].astype('category').cat.codes

# --- Define features ---
X = df[['Density (Plants)', 'Rainfall intensity (mm/h)']].values
T = df['Treatments_encoded'].values

# --- List of outcomes ---
outcomes = ['Soil loss (ton/h)', 'Runoff (mm)', 'TN (kg/h)', 'TP (kg/h)']

# --- Storage for results ---
results = {}

# --- Run Causal Forest for each outcome ---
for outcome in outcomes:
    Y = df[outcome].values
    model = CausalForestDML(random_state=42)
    model.fit(Y, T, X=X, W=None)  # no cluster_ids in new versions
    
    ate = model.ate(X)
    ite = model.effect(X)
    feature_importances = model.feature_importances_
    
    results[outcome] = {
        'ATE': ate,
        'ITE': ite,
        'Feature Importances': feature_importances
    }

# --- Display results ---
for outcome, res in results.items():
    print(f"\nOutcome: {outcome}")
    print("Average Treatment Effect (ATE):", res['ATE'])
    print("First 5 individual treatment effects (ITE):", res['ITE'][:5])
    print("Feature importances:", res['Feature Importances'])
