import os
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

MONTHS = {
    'Januari':1,'Februari':2,'Maret':3,'April':4,'Mei':5,'Juni':6,
    'Juli':7,'Agustus':8,'September':9,'Oktober':10,'November':11,'Desember':12
}
MONTH_NAMES = {v:k for k,v in MONTHS.items()}
FEATURES = ['Harga','Stok','bulan','tahun','lag_1','lag_2','lag_3']


def parse_indonesian_period(value):
    if pd.isna(value):
        return pd.NaT
    if isinstance(value, (pd.Timestamp,)): return value
    text = str(value).strip()
    parts = text.split()
    if len(parts) != 2 or parts[0] not in MONTHS:
        return pd.to_datetime(text, errors='coerce')
    try:
        return pd.Timestamp(year=int(parts[1]), month=MONTHS[parts[0]], day=1)
    except (TypeError, ValueError):
        return pd.NaT


def prepare_features(df):
    d = df.copy()
    if 'periode_date' in d.columns:
        d['periode_date'] = pd.to_datetime(d['periode_date'], errors='coerce')
    else:
        d['periode_date'] = d['periode'].map(parse_indonesian_period)
    d = d.sort_values(['Produk','periode_date']).reset_index(drop=True)
    d['bulan'] = d['periode_date'].dt.month
    d['tahun'] = d['periode_date'].dt.year
    g = d.groupby('Produk', group_keys=False)
    d['lag_1'] = g['Terjual'].shift(1)
    d['lag_2'] = g['Terjual'].shift(2)
    d['lag_3'] = g['Terjual'].shift(3)
    return d


def _model(n_estimators=300, random_state=42):
    return RandomForestRegressor(
        n_estimators=n_estimators,
        max_depth=None,
        min_samples_leaf=1,
        random_state=random_state,
        n_jobs=-1,
    )


def train_evaluate(df, n_estimators=300, random_state=42):
    d = prepare_features(df).dropna(subset=FEATURES + ['Terjual']).copy()
    train = d[d['periode_date'] < pd.Timestamp('2026-01-01')].copy()
    test = d[d['periode_date'] >= pd.Timestamp('2026-01-01')].copy()
    if train.empty or test.empty:
        raise ValueError('Data training/testing tidak mencukupi untuk evaluasi.')
    model = _model(n_estimators, random_state)
    model.fit(train[FEATURES], train['Terjual'])
    test['prediksi'] = np.maximum(0, model.predict(test[FEATURES]))
    metrics = {
        'MAE': mean_absolute_error(test['Terjual'], test['prediksi']),
        'RMSE': mean_squared_error(test['Terjual'], test['prediksi']) ** 0.5,
        'R2': r2_score(test['Terjual'], test['prediksi']),
    }
    return model, d, train, test, metrics


def fit_full(df, n_estimators=300, random_state=42):
    d = prepare_features(df).dropna(subset=FEATURES + ['Terjual']).copy()
    if d.empty:
        raise ValueError('Data dengan fitur lengkap tidak tersedia.')
    model = _model(n_estimators, random_state)
    model.fit(d[FEATURES], d['Terjual'])
    return model, d


def next_month_forecast(df, model=None):
    d = prepare_features(df)
    if model is None:
        model, _ = fit_full(df)
    latest_date = d['periode_date'].max()
    target_date = (latest_date + pd.offsets.MonthBegin(1)).normalize()
    rows = []
    for product, g in d.groupby('Produk'):
        g = g.sort_values('periode_date').dropna(subset=['Terjual'])
        if len(g) < 3:
            continue
        latest = g.iloc[-1]
        row = {
            'Produk': product,
            'Kategori': latest['Kategori'],
            'Periode Prediksi': target_date.strftime('%Y-%m'),
            'Harga': latest['Harga'],
            'Stok': latest['Stok'],
            'bulan': target_date.month,
            'tahun': target_date.year,
            'lag_1': latest['Terjual'],
            'lag_2': g.iloc[-2]['Terjual'],
            'lag_3': g.iloc[-3]['Terjual'],
        }
        row['hasil_prediksi'] = max(0.0, float(model.predict(pd.DataFrame([row])[FEATURES])[0]))
        rows.append(row)
    return pd.DataFrame(rows), target_date


def save_model(model, path='model/random_forest.pkl'):
    os.makedirs(os.path.dirname(path) or '.', exist_ok=True)
    joblib.dump(model, path)