from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder

class DurationTransform(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self

    def transform(self, X):
        import pandas as pd
        if isinstance(X, pd.DataFrame):
            df = X.copy()
        else:
            df = pd.DataFrame(X, columns=["transactiondate", "membershipexpiredate"])
        df["transactiondate"] = pd.to_datetime(df["transactiondate"])
        df["membershipexpiredate"] = pd.to_datetime(df["membershipexpiredate"])
        result = (df["membershipexpiredate"] - df["transactiondate"]).dt.days
        return result.values.reshape(-1, 1)

gen_encoding = ColumnTransformer([
    ("gender", OneHotEncoder(), [1])
], remainder='passthrough')

sub_time = ColumnTransformer([
    ("duration_in_days", DurationTransform(), [8, 9])
], remainder='passthrough')

pipe = Pipeline([
    ('gen_encoding', gen_encoding),
    ('sub_time', sub_time)
])

from sklearn.utils import resample

# Chia nhóm khách hàng trung thành và rời bỏ
zeros = dbtrain[dbtrain['ischurn'] == 0]   # 9,354 non-churners
ones = dbtrain[dbtrain['ischurn'] == 1]    # 646 churners

# Giảm số lượng khách trung thành để cân bằng với churners
zeros_undersampled = resample(zeros, replace=False, 
                              n_samples=len(ones), random_state=42)

# Kết hợp và xáo trộn dữ liệu
dbtrain_balanced = pd.concat([zeros_undersampled, ones])
dbtrain_balanced = dbtrain_balanced.sample(frac=1, random_state=42).reset_index(drop=True)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from datetime import date

app = FastAPI(title="Customer Churn Prediction",
              description="API ML sản xuất cho dự đoán churn",
              version="1.0.0")

class DataVal(BaseModel):
    userid: Optional[int] = None
    city: int
    gender: str
    registeredvia: int
    paymentmethodid: int
    paymentplandays: int
    actualamountpaid: int
    isautorenew: int
    transactiondate: date
    membershipexpiredate: date

@app.post("/predict")
def predict(data: DataVal):
    # Validate và xử lý dữ liệu
    # Biến đổi dữ liệu qua pipeline, dự đoán, lưu kết quả
    try:
        # Chuyển dữ liệu đầu vào thành dạng list
        pipedata = [[
            data.city, data.gender, data.registeredvia, data.paymentmethodid,
            data.paymentplandays, data.actualamountpaid, data.isautorenew,
            data.transactiondate, data.membershipexpiredate
        ]]
        transformed = pipe.transform(pipedata)
        df_transformed = pd.DataFrame(transformed, columns=[
            "durationofsubscription", "female", "male", "city", "registeredvia", 
            "paymentmethodid", "paymentplandays", "actualamountpaid", "isautorenew"
        ])

        prediction = model.predict(df_transformed)
        result = {data.userid or generate_userid(): df_transformed.iloc[0].to_dict()}
        result[data.userid]["prediction"] = int(prediction[0])
        
        save_prediction(result)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Dự đoán thất bại: {str(e)}")

import json
import os

def valid_user(user: Optional[int]) -> int:
    """Generate or validate user ID with persistent storage"""
    if user is None:
        with open("data/users.json", "r") as f:
            data = json.load(f)
        max_user = max(data) if data else 0
        user = max_user + 1
        data.append(user)
        with open("data/users.json", "w") as f:
            json.dump(data, f, indent=2)
    else:
        with open("data/users.json", "r") as f:
            data = json.load(f)
        if user not in data:
            data.append(user)
            with open("data/users.json", "w") as f:
                json.dump(data, f, indent=2)
    return user

def save_prediction(result: dict):
    """Save prediction result with user data"""
    json_path = "data/userdata.json"
    if os.path.exists(json_path):
        with open(json_path, "r") as f:
            json_file = json.load(f)
    else:
        json_file = {}
    json_file.update(result)
    with open(json_path, "w") as f:
        json.dump(json_file, f, indent=2)

import cloudpickle

# Lưu mô hình và pipeline
with open("model/model.pickle", "wb") as f:
    cloudpickle.dump(adaboost_model, f)

with open("model/pipe.pickle", "wb") as f:
    cloudpickle.dump(pipe, f)

# Tải mô hình và pipeline phục vụ sản xuất
with open("model/model.pickle", "rb") as f:
    model = cloudpickle.load(f)

with open("model/pipe.pickle", "rb") as f:
    pipe = cloudpickle.load(f)