import sys
import os
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# Add the directory containing the 'ellice' package to sys.path
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import ellice
from ellice.utils.helpers import compute_ellipsoid

def create_simple_data():
    # 2 continuous features, 1 ordinal feature (0,1,2)
    X = np.random.randn(100, 2)
    ordinal = np.random.randint(0, 3, 100)
    X = np.c_[X, ordinal]
    y = (X[:, 0] + X[:, 1] + X[:, 2] > 0).astype(int)
    
    df = pd.DataFrame(X, columns=['cont1', 'cont2', 'ord1'])
    return df, pd.Series(y, name='target')

def test_ellice_units():
    print("Running Unit Tests...")
    
    # Setup Data & Model
    df, y = create_simple_data()
    X_train, X_test, y_train, y_test = train_test_split(df, y, test_size=0.2, random_state=42)
    
    model = LogisticRegression()
    model.fit(X_train, y_train)
    
    full_df = X_train.copy()
    full_df['target'] = y_train
    data = ellice.Data(full_df, target_column='target')
    exp = ellice.Explainer(model, data, backend='sklearn')
    
    query = X_test.iloc[0]
    target_class = 1 - model.predict([query])[0]
    
    print(f"Query: {query.values}, Target Class: {target_class}")

    # Test 1: Continuous Generator - Basic
    print("\n[Test 1] Continuous Generator - Basic")
    # max_iterations passed via optimization_params
    cf = exp.generate_counterfactuals(
        query, target_class=target_class, method='continuous', 
        progress_bar=False,
        optimization_params={'max_iterations': 100}
    )
    if not cf.empty:
        print("PASS: Generated CF")
    else:
        print("FAIL: No CF generated")

    # Test 2: Continuous Generator - Early Stopping
    print("\n[Test 2] Early Stopping (Robustness)")
    # Force highly robust result requirement
    cf = exp.generate_counterfactuals(
        query, target_class=target_class, method='continuous', 
        progress_bar=False,
        optimization_params={'max_iterations': 500}
    )
    if not cf.empty:
        # Check if robust logit > 0 (implied by successful generation, but let's verify metric)
        # We can check the 'worst_case_prob_target' if return_probs=True, but here we just check existence
        print("PASS: Early stopping worked (or max iters reached)")
    else:
        print("FAIL: No CF generated")
    
    # Test 3: Allowed Values (Ordinal)
    print("\n[Test 3] Allowed Values (Ordinal)")
    allowed = {'ord1': [0.0, 1.0, 2.0]}
    cf = exp.generate_counterfactuals(
        query, target_class=target_class, method='continuous',
        allowed_values=allowed, progress_bar=False,
        optimization_params={'max_iterations': 100}
    )
    if not cf.empty:
        val = cf.iloc[0]['ord1']
        if any(np.isclose(val, v, atol=1e-4) for v in allowed['ord1']):
            print(f"PASS: Ordinal value {val} is allowed")
        else:
            print(f"FAIL: Ordinal value {val} NOT allowed")

    # Test 4: Discrete Generator
    print("\n[Test 4] Discrete Generator")
    cf = exp.generate_counterfactuals(
        query, target_class=target_class, method='discrete', progress_bar=False
    )
    if not cf.empty:
        print("PASS: Discrete CF found")
    else:
        print("INFO: No discrete candidate found (might be expected for sparse data)")

if __name__ == "__main__":
    test_ellice_units()
