import torch
import numpy as np
from transformers import pipeline
from sklearn.preprocessing import MinMaxScaler

def predict_umkm(data, prediction_length=7, safety_stock=10):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    try:
        # ===== 1. Persiapan Data =====
        scaler = MinMaxScaler()
        scaled_demand = scaler.fit_transform(data[['demand']]).flatten().tolist()
        
        # ===== 2. Prediksi dengan Chronos-T5 =====
        chronos = pipeline(
            "text-generation",
            model="amazon/chronos-t5-small",
            device=device
        )
        
        # Format prompt khusus time series
        prompt = f"""
        [INSTRUCTION]
        Berikan rekomendasi manajemen inventori untuk {prediction_length} hari ke depan:
        - Data historis demand: {scaled_demand[-100:]}  # Ambil 100 data terakhir
        - Stok saat ini: {data['supply'].iloc[-1]}
        - Safety stock: {safety_stock}
        
        [FORMAT]
        1 kalimat dalam Bahasa Indonesia dengan angka konkret
        Estimasi ROI dalam persentase
        [/FORMAT]
        """
        
        # Generate rekomendasi
        response = chronos(
            prompt,
            max_new_tokens=100,
            temperature=0.7
        )[0]['generated_text']
        
        # Ekstrak teks rekomendasi
        rekomendasi = response.split("[/FORMAT]")[-1].strip()
        
        # ===== 3. Simulasi ROI ===== (Contoh Sederhana)
        avg_demand = np.mean(data['demand'][-30:])
        roi = round((avg_demand * 0.2) / (data['supply'].iloc[-1] + safety_stock) * 100, 1)
        
        return {
            "rekomendasi": rekomendasi,
            "prediksi": [avg_demand] * prediction_length,  # Contoh prediksi
            "roi": roi,
            "stok_optimal": int(avg_demand * 1.2) + safety_stock,
            "anomali": len(data) - len(clean_data(data))
        }
        
    except Exception as e:
        return {"error": str(e)}