from flask_restx import Resource, Namespace
from flask import request, jsonify
import os
import json
from database.database import DB
import torch
from torchvision.transforms import ToTensor
from datetime import datetime
from model.AttentiveRNN import AttentiveRNN
from model.Classifier import Resnet as Classifier
from subfuction.image_crop import crop_image
import numpy as np
import cv2

db = DB()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# pre-loading models
arnn = AttentiveRNN(6, 3, 2)
arnn.eval()
arnn.load_state_dict(torch.load("weights/ARNN_trained_weight_6_3_2.pt"))
arnn.to(device=device)

classifier = Classifier()
classifier.eval()
classifier.load_state_dict(torch.load("weights/Classifier_512.pt"))
classifier.to(device=device)

tf_toTensor = ToTensor()
crop_size = (512, 512)
start_point = (750, 450)
root_dir = os.getcwd()

Action = Namespace(
    name="Action",
    description="노드 분석을 위해 사용하는 api.",
)


# @Action.route('/image_summit')
# class fileUpload(Resource):
#     @Action.doc(responses={200: 'Success'})
#     @Action.doc(responses={500: 'Register Failed'})
#     def post(self):
#         if request.method == 'POST':
#             f = request.files['file']
#             f.save(secure_filename(f.filename))
#             return {
#                 'save': 'done'  # str으로 반환하여 return
#             }, 200


@Action.route('/image_anal')
class fileUpload(Resource):
    @Action.doc(responses={200: 'Success'})
    @Action.doc(responses={500: 'Register Failed'})
    def post(self):
        # Extracting JSON data
        json_data = request.form.get('data')
        if not json_data:
            return jsonify({"message": "Missing JSON data"}), 400
        data = json.loads(json_data)

        lat = float(data['gps_x'])
        lon = float(data['gps_y'])
        filename = data['filename']
        file_type = data['file_type']

        uploaded_file = request.files.get('file')
        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
        image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)

        image = crop_image(image, crop_size, start_point)

        image_tensor = tf_toTensor(image)
        image_tensor = image_tensor.unsqueeze(0)
        image_tensor = image_tensor.to(device)
        with torch.no_grad():
            image_arnn = arnn(image_tensor)
            image_tensor.cpu()
            del image_tensor
            result = classifier(image_arnn['x'])
            image_arnn['x'].cpu()
            del image_arnn

        result = result.to("cpu")
        _, predicted = torch.max(result.data, 1)
        del result
        if predicted == 0:
            rain = False
        else:  # elif result == 1
            rain = True
        user_id = 'test'
        action_success = True
        action_id = 'test'
        db.db_add_action(action_id, lat, lon, user_id, action_success)
        return {
            'node': (lat, lon),
            'rain': rain,
        }, 200


@Action.route('/action_display')
class fileUpload(Resource):
    @Action.doc(responses={200: 'Success'})
    @Action.doc(responses={500: 'Register Failed'})
    def get(self):
        if request.method == 'GET':
            db = DB()
            value = db.db_display_action()
            return {
                'report': list(value)
            }, 200