import xgboost as xgb
import pandas as pd
from datetime import datetime, timedelta
import os

# Extended dummy dataset
data = pd.DataFrame([
    {"last_premium_paid_date": (datetime.now() - timedelta(days=60)).strftime('%Y-%m-%d'), "payment_mode": "Annual", "policy_term": 15, "policy_age": 3, "risk": 1},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=10)).strftime('%Y-%m-%d'), "payment_mode": "Monthly", "policy_term": 20, "policy_age": 2, "risk": 0},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=400)).strftime('%Y-%m-%d'), "payment_mode": "Quarterly", "policy_term": 25, "policy_age": 5, "risk": 1},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=700)).strftime('%Y-%m-%d'), "payment_mode": "Semi-Annual", "policy_term": 10, "policy_age": 8, "risk": 1},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=90)).strftime('%Y-%m-%d'), "payment_mode": "Annual", "policy_term": 12, "policy_age": 4, "risk": 0},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=30)).strftime('%Y-%m-%d'), "payment_mode": "Monthly", "policy_term": 20, "policy_age": 1, "risk": 0},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=300)).strftime('%Y-%m-%d'), "payment_mode": "Annual", "policy_term": 15, "policy_age": 3, "risk": 1},
    {"last_premium_paid_date": (datetime.now() - timedelta(days=180)).strftime('%Y-%m-%d'), "payment_mode": "Quarterly", "policy_term": 18, "policy_age": 6, "risk": 0},
])

def encode_payment_mode(mode):
    return {"Annual": 0, "Semi-Annual": 1, "Quarterly": 2, "Monthly": 3}.get(mode, -1)

def calculate_months_since(date_str):
    try:
        delta = datetime.now() - datetime.strptime(date_str, "%Y-%m-%d")
        return delta.days // 30
    except:
        return 0

data["months_since_last_payment"] = data["last_premium_paid_date"].apply(calculate_months_since)
data["payment_mode_encoded"] = data["payment_mode"].apply(encode_payment_mode)
X = data[["months_since_last_payment", "payment_mode_encoded", "policy_term", "policy_age"]]
y = data["risk"]

model = xgb.XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X, y)

os.makedirs("model", exist_ok=True)
model.save_model("model/xgb_model.json")