import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

def generate_synthetic_data(n_samples: int = 1000, random_state: int = 42) -> pd.DataFrame:
    """
    Generates robust synthetic customer churn data.
    
    Features:
    - age: customer age (int, years)
    - monthly_charges: monthly subscription charge (float, USD)
    - contract_length: months of commitment (int: 1, 12, 24)
    - support_calls: number of customer support calls (int)
    - tech_support: whether customer has tech support (str: 'yes', 'no')
    - churn: label (int: 0, 1)
    """
    np.random.seed(random_state)
    
    age = np.random.randint(18, 70, size=n_samples)
    monthly_charges = np.random.uniform(20.0, 120.0, size=n_samples)
    contract_length = np.random.choice([1, 12, 24], size=n_samples, p=[0.4, 0.4, 0.2])
    support_calls = np.random.poisson(lam=1.5, size=n_samples)
    tech_support = np.random.choice(['yes', 'no'], size=n_samples, p=[0.3, 0.7])
    
    # Generate target probability with some clear mathematical relations (churn triggers)
    # Higher support calls -> higher churn
    # Contract 1 month -> higher churn
    # Higher charges -> higher churn
    # Tech support 'yes' -> lower churn
    churn_logit = (
        0.5 * support_calls
        - 0.02 * age
        + 0.01 * monthly_charges
        - 1.2 * (contract_length > 1)
        - 0.8 * (tech_support == 'yes')
        - 0.5 # bias
    )
    
    churn_prob = 1 / (1 + np.exp(-churn_logit))
    churn = (np.random.rand(n_samples) < churn_prob).astype(int)
    
    df = pd.DataFrame({
        'age': age,
        'monthly_charges': monthly_charges,
        'contract_length': contract_length,
        'support_calls': support_calls,
        'tech_support': tech_support,
        'churn': churn
    })
    
    return df

class Preprocessor:
    """
    Handles preprocessing: scaling numerical columns and converting categorical values.
    Uses scikit-learn standard Scaler internally and handles one-hot encoding manual/safe mapping.
    """
    def __init__(self):
        self.scaler = StandardScaler()
        self.numerical_cols = ['age', 'monthly_charges', 'contract_length', 'support_calls']
        self.categorical_cols = ['tech_support']
        self.is_fitted = False
        
    def fit(self, df: pd.DataFrame):
        # Fit scaler on numericals
        self.scaler.fit(df[self.numerical_cols])
        self.is_fitted = True
        return self
        
    def transform(self, df: pd.DataFrame) -> np.ndarray:
        if not self.is_fitted:
            raise ValueError("Preprocessor has not been fitted yet!")
            
        # Scale numericals
        num_scaled = self.scaler.transform(df[self.numerical_cols])
        
        # Binary encode/one-hot the tech_support (yes=1, no=0)
        tech_support_encoded = (df['tech_support'] == 'yes').astype(float).values.reshape(-1, 1)
        
        # Combine features
        features = np.hstack([num_scaled, tech_support_encoded])
        return features

    def fit_transform(self, df: pd.DataFrame) -> np.ndarray:
        return self.fit(df).transform(df)

def prepare_data(df: pd.DataFrame, test_size: float = 0.2, random_state: int = 42):
    """
    Splits the data, preprocesses features and returns train/test splits.
    """
    X = df.drop(columns=['churn'])
    y = df['churn']
    
    X_train_raw, X_test_raw, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )
    
    preprocessor = Preprocessor()
    X_train = preprocessor.fit_transform(X_train_raw)
    X_test = preprocessor.transform(X_test_raw)
    
    return X_train, X_test, y_train.values, y_test.values, preprocessor