import os
import json
import numpy as np
import torch
import torch.nn as nn
import joblib
from flask import Flask, request, jsonify, render_template

# ─────────────────────────────────────────
#  Model definition  (must match training)
# ─────────────────────────────────────────
class LSTMModel(nn.Module):
    def __init__(self, input_size=10, hidden_size=64, num_layers=2, output_size=1):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc   = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out, _ = self.lstm(x)          # out: (batch, seq_len, hidden)
        out    = self.fc(out[:, -1, :])  # last timestep → (batch, 1)
        return out

# ─────────────────────────────────────────
#  Load model & scaler once at startup
# ─────────────────────────────────────────
BASE = os.path.dirname(__file__)

model = LSTMModel()
state = torch.load(os.path.join(BASE, "model_LSTM.pth"), map_location="cpu")
model.load_state_dict(state)
model.eval()

scaler = joblib.load(os.path.join(BASE, "scaler.joblib"))

FEATURE_NAMES = [
    "Active_Energy_Delivered_Received",
    "Current_Phase_Average",
    "Active_Power",
    "Wind_Speed",
    "Weather_Temperature_Celsius",
    "Weather_Relative_Humidity",
    "Global_Horizontal_Radiation",
    "Diffuse_Horizontal_Radiation",
    "Wind_Direction",
    "Weather_Daily_Rainfall",
]

# ─────────────────────────────────────────
#  Flask app
# ─────────────────────────────────────────
app = Flask(__name__)


@app.route("/")
def index():
    return render_template("index.html", features=FEATURE_NAMES)


@app.route("/predict", methods=["POST"])
def predict():
    try:
        data = request.get_json(force=True)

        # Expect: { "rows": [[f1,f2,...,f10], [f1,f2,...,f10], ...] }
        rows = data.get("rows", [])
        if not rows:
            return jsonify({"error": "No input rows provided."}), 400

        # Validate shape
        for i, row in enumerate(rows):
            if len(row) != 10:
                return jsonify({"error": f"Row {i+1} must have exactly 10 values."}), 400

        X = np.array(rows, dtype=np.float32)          # (seq_len, 10)
        X_scaled = scaler.transform(X)                 # scale each timestep
        tensor = torch.tensor(X_scaled).unsqueeze(0)  # (1, seq_len, 10)

        with torch.no_grad():
            pred = model(tensor).item()

        return jsonify({
            "prediction": round(pred, 6),
            "unit": "Active Energy (scaled output)",
            "seq_len": len(rows),
        })

    except Exception as exc:
        return jsonify({"error": str(exc)}), 500


@app.route("/health")
def health():
    return jsonify({"status": "ok", "model": "LSTM Solar Predictor"})


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    debug = os.environ.get("FLASK_DEBUG", "0") == "1"
    app.run(host="0.0.0.0", port=port, debug=debug)