import pandas as pd
import torch
import numpy as np
from modules.config import SEQ_LEN
from datetime import datetime, timedelta


def create_sequences(data, seq_len):
    xs = []
    for i in range(seq_len, len(data)):
        x = data[i - seq_len:i]
        xs.append(x)
    return np.array(xs)

def infer_aqi(df, location_id, datetime_str, scaler_list, model_point):
    df_loc = df[(df["ID Vị Trí"] == location_id)].copy()

    # Chuyển đổi thời gian
    df_loc["Datetime"] = pd.to_datetime(df_loc["Datetime"], format="%H:%M %d/%m/%Y", errors='coerce')

    # Lọc theo thời gian <= thời điểm được chọn (lấy dữ liệu trước thời điểm này)
    try:
        target_time = pd.to_datetime(datetime_str, format="%H:%M %d/%m/%Y")
    except Exception as e:
        print(f"❌ Định dạng thời gian sai. Định dạng đúng: HH:MM dd/mm/yyyy\nChi tiết: {e}")
        return 0.0

    df_loc = df_loc[df_loc["Datetime"] <= target_time]
    df_loc = df_loc.sort_values("Datetime")

    if len(df_loc) < SEQ_LEN:
        print("❌ Không đủ dữ liệu trước thời điểm này để dự đoán")
        return 0.0

    # Xử lý dữ liệu tương tự khi train
    df_loc["AQI_PM2.5"] = df_loc["AQI_PM2.5"].astype(str).str.replace(",", ".")
    df_loc["AQI_PM2.5"] = pd.to_numeric(df_loc["AQI_PM2.5"], errors="coerce")

    data_input = df_loc[["AQI_PM2.5", "Vĩ độ", "Kinh độ"]].values
    scaler = scaler_list[location_id]
    data_scaled = scaler.transform(data_input)

    # Tạo chuỗi đầu vào
    seq_input = create_sequences(data_scaled, SEQ_LEN)
    input_tensor = torch.FloatTensor(seq_input[-1:])  # lấy chuỗi gần nhất

    model = model_point[location_id]
    model.eval()
    with torch.no_grad():
        pred_scaled = model(input_tensor).numpy().squeeze()

    # Đảo chuẩn hóa
    pred_extended = np.zeros((1, 3))
    pred_extended[:, 0] = pred_scaled
    pred_original = scaler.inverse_transform(pred_extended)[0, 0]
    
    return pred_original

def infer_and_append_aqi(df, location_id, datetime_str, scaler_list, model_point):
    # Bước 1: Thực hiện dự đoán như cũ
    pred_list = []
    for location_id in range(0, len(scaler_list)):
        pred_original = infer_aqi(df, location_id, datetime_str, scaler_list, model_point)
        
        # Bước 2: Tạo bản ghi mới từ kết quả dự đoán
        original_time = datetime.strptime(datetime_str, "%H:%M %d/%m/%Y")
        one_hour_later = original_time + timedelta(hours=1)
        formatted_time = one_hour_later.strftime("%H:%M %d/%m/%Y")
        new_record = {
            "ID Vị Trí": location_id,
            "Datetime": formatted_time,
            "AQI_PM2.5": pred_original,
            "Tên": df[df["ID Vị Trí"] == location_id]["Tên"].iloc[0],
            "Vĩ độ": df[df["ID Vị Trí"] == location_id]["Vĩ độ"].iloc[0],
            "Kinh độ": df[df["ID Vị Trí"] == location_id]["Kinh độ"].iloc[0],
            "is_predicted": True  # Đánh dấu là dữ liệu dự đoán
        }
        
        # Bước 3: Thêm vào DataFrame
        # Sử dụng biến toàn cục hoặc truyền df vào theo cách khác
        df = pd.concat([df, pd.DataFrame([new_record])], ignore_index=True)
        pred_list.append(pred_original)
    return pred_list, df.copy()