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

def get_data_from_postgres(conn_string):
    """
    Fetch dating app data from PostgreSQL and prepare for model training
    """
    try:
        print(f"Connecting to PostgreSQL database: {conn_string}")

        engine = create_engine(conn_string)
        print("Connection successful", engine)

        query = """
        SELECT
            id,
            hobies_matched,
            is_job_matched,
            is_edu_matched,
            is_religion_match,
            is_interested_in_match,
            profile_completion,
            no_of_photos,
            miles_away,
            user_id,
            is_liked,
            age
        FROM
            ml_data
        """

        df = pd.read_sql(query, engine)

        print(f"Successfully fetched {len(df)} records from PostgreSQL")

        bool_cols = ['is_job_matched', 'is_edu_matched', 'is_religion_match',
                      'is_interested_in_match', 'is_liked']
        for col in bool_cols:
            if df[col].dtype == bool:
                df[col] = df[col].astype(int)

        return df

    except Exception as e:
        print(f"Error connecting to PostgreSQL database: {e}")
        return None

def preprocess_data(df, test_size=0.20, random_state=42):
    """
    Preprocess the dating app data for logistic regression
    """
    data = df.copy()

    # data['compatibility_score'] = (
    #     data['hobies_matched'] * 0.3 +
    #     data['is_job_matched'] * 0.1 +
    #     data['is_edu_matched'] * 0.1 +
    #     data['is_religion_match'] * 0.2 +
    #     data['is_interested_in_match'] * 0.3
    # )

    # data['miles_away_log'] = np.log1p(data['miles_away'])

    # data['profile_quality'] = (data['profile_completion'] * 0.7 +
    #                           data['no_of_photos'] * 30 * 0.3)

    # data['interest_x_photos'] = data['is_interested_in_match'] * data['no_of_photos']
    # data['hobbies_x_religion'] = data['hobies_matched'] * data['is_religion_match']

    exclude_cols = ['id', 'user_id', 'is_liked']
    feature_cols = [col for col in data.columns if col not in exclude_cols]
    print(f"Feature columns: {feature_cols}")

    X = data[feature_cols].values
    y = data['is_liked'].values

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state, stratify=y
    )

    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    return X_train, X_test, y_train, y_test, scaler

def transform_new_data(data, scaler):
    """
    Transform new data for prediction using the fitted scaler
    """
    # Convert to numpy array if it's a DataFrame
    if isinstance(data, pd.DataFrame):
        data = data.values

    # Apply the same scaling used during training
    scaled_data = scaler.transform(data)

    return scaled_data