import gradio as gr
import torch
import numpy as np
import json
from Attn_conv import Attn_conv  # 替换为你的模型路径

# 加载模型
model = Attn_conv(n_class=4)
model.load_state_dict(torch.load("./Attn_conv.pth", weights_only=True,map_location=torch.device('cpu')))
model.eval()

# 推理函数
def predict(data):
    # 读取上传的 JSON 文件
    try:
        # with open(file.name, "r") as f:
        #     data = json.load(f)

        channels = ["AccelX", "AccelY", "AccelZ", "GyroX", "GyroY", "GyroZ"]

        # 转换为 NumPy 数组，形状为 (6, sequence_length)
        imu_data = np.array([data[channel] for channel in channels])

        #Process data
        len_data = imu_data.shape[1]
        window_size = 30
        overlap = 15
        step = window_size - overlap
        predictions = [0,0,0,0]
        

        for start in range(0,len_data,step):
            end = start + window_size
            if end > len_data:
                # Zero padding to last window
                window = np.pad(imu_data[:, start:], ((0, 0), (0, end - len_data)), mode='constant')
            else:
                window = imu_data[:, start:end]
            
            input_tensor = torch.tensor(window, dtype=torch.float32).unsqueeze(0)

            with torch.no_grad():
                output = model(input_tensor) 
                probabilities = torch.softmax(output, dim=-1).cpu().numpy()
                predicted_class = np.argmax(probabilities)
                #print(probabilities)
                if np.max(probabilities) > 0.70:    #Threshold
                    predictions[predicted_class] += 1
        #取眾數
        max_count = max(predictions)
        predictions = [x if x == max_count else 0 for x in predictions]

        #Result
        # result = (
        #     f"bicep: {predictions[0]} | abs: {predictions[1]} | "
        #     f"chess: {predictions[2]} | legs: {predictions[3]}"
        # )
        class_names = ["dumbbells", "situps", "pushups", "squats"]
        if max_count == 0:
            # 沒有任何類別達到閾值
            result_json = {
                "type": None,
                "reps": 0
            }
        else:
            # 找出哪個類別擁有 max_count
            # 如果有多個類別並列最高，這裡僅選第一個
            max_index = predictions.index(max_count)
            result_json = {
                "type": class_names[max_index],
                "reps": max_count
            }

        return result_json
    except Exception as e:
        return {"error": str(e)}

# 定义 Gradio 界面
iface = gr.Interface(
    fn=predict,
    inputs=gr.JSON(label="json file"),
    outputs=gr.JSON(label="Prediction Result"),  
    title="Burnsync classifier",
    description="Upload a JSON file with 6-channel IMU time series data to predict the action class."
)

# 启动应用
iface.launch()