import cv2
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.pipeline import make_pipeline
from sklearn.base import BaseEstimator, TransformerMixin
import cv2
from keras.models import Model
from keras.layers import Input
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Concatenate
from keras.utils import plot_model
import tensorflow as tf
from keras.layers import GlobalAveragePooling2D, GlobalMaxPooling2D, Reshape, Dense, Input, Flatten, Dropout
from keras.layers import Activation, Concatenate, Conv2D, Multiply
from keras.applications.inception_v3 import InceptionV3
import keras
from keras.applications.inception_v3 import preprocess_input
import tensorflow as tf
from keras import layers
from sklearn.model_selection import train_test_split
import pandas as pd

import numpy as np

class RoiExtractor:
    def __init__(self):
        # Model initialization
        self.model = build_image_model()
        self.model.load_weights("https://drive.google.com/file/d/1JqIQHYvW8q5DMjOpcVsfkwxaOZP2CmnL/view?usp=sharing")
        self.heatmap_model = Model(inputs=self.model.inputs, outputs=self.model.layers[-3].output)

        # Outputs
        self.heatmap_1 = None
        self.carpal_img = None
        
        self.heatmap_2 = None
        self.metacarpal_img = None
        
        self.img = None
        self.masked_img = None

    def generate_heatmap(self, img):
        
        # Preprocessing
        i = preprocess_input(img)
        preprocessed_img = np.expand_dims(i, axis=0)

        cbam_output = self.heatmap_model.predict(preprocessed_img)[0]

        #Heatmap generation
        heatmap = np.sum(cbam_output, axis=-1)
        heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap))
        upsampled_heatmap = cv2.resize(heatmap, (299, 299))
        return upsampled_heatmap


    def get_bounding_box(self, heatmap, threshold=0.7):
        binary_mask = (heatmap > threshold).astype(np.uint8)
        non_zero_indices = np.nonzero(binary_mask)
        if len(non_zero_indices[0]) == 0 or len(non_zero_indices[1]) == 0:
            return None
        min_row, min_col = np.min(non_zero_indices[0]), np.min(non_zero_indices[1])
        max_row, max_col = np.max(non_zero_indices[0]), np.max(non_zero_indices[1])

        bounding_box ={
            'min_row': min_row,
            'min_col': min_col,
            'max_row': max_row,
            'max_col': max_col
        }
        return bounding_box


    def crop_image_by_bounding_box(self, image, bounding_box):
        min_row, min_col = bounding_box['min_row'], bounding_box['min_col']
        max_row, max_col = bounding_box['max_row'], bounding_box['max_col']
        cropped_image = image[min_row:max_row + 1, min_col:max_col + 1, :]
        cropped_resized = cv2.resize(cropped_image, (224, 224))
        return cropped_resized


    def apply_black_rectangle_by_bounding_box(self, image, bounding_box):
        min_row, min_col = bounding_box['min_row'], bounding_box['min_col']
        max_row, max_col = 299, bounding_box['max_col']
        test_img = image.copy()
        test_img = cv2.rectangle(test_img, (min_col, min_row), (max_col, max_row), (0, 0, 0), thickness=cv2.FILLED)
        return test_img

    
    def process_img(self, img): 
        self.img = img
        self.heatmap_1 = self.generate_heatmap(img)
        bounding_box_1 = self.get_bounding_box(self.heatmap_1)
        self.carpal_img = self.crop_image_by_bounding_box(img, bounding_box_1)
        self.masked_img = self.apply_black_rectangle_by_bounding_box(img, bounding_box_1)
        self.heatmap_2 = self.generate_heatmap(self.masked_img)
        bounding_box_2 = self.get_bounding_box(self.heatmap_2)
        self.metacarpal_img = self.crop_image_by_bounding_box(img, bounding_box_2)
    

# Channel Attention Module
def channel_attention_module(x, ratio=8):

    batch,_,_,channel=x.shape
    # shared layers
    l1 = Dense(channel//ratio, activation="relu", use_bias=False)
    l2 = Dense(channel, use_bias= False)

    x1 = GlobalAveragePooling2D()(x)
    x1 = Flatten()(x1)
    x1 = l1(x1)
    x1 = l2(x1)

    x2 = GlobalMaxPooling2D()(x)
    x2 = Flatten()(x2)
    x2 = l1(x2)
    x2 = l2(x2)

    feats = x1 + x2
    feats = Activation("sigmoid")(feats)
    feats = Multiply()([x,feats])

    return feats

# spatical attention module

def spatial_attention_module(x):
    # Average Pooling
    x1 = tf.reduce_mean(x,axis = -1)
    x1 = tf.expand_dims(x1,axis = -1)

    # max pooling
    x2 = tf.reduce_max(x, axis = -1)
    x2 = tf.expand_dims(x2,axis=-1)

    feats = Concatenate()([x1,x2])

    feats = Conv2D(1,kernel_size=7, padding="same",activation="softmax")(feats)
    feats = Multiply()([x,feats])

    return feats

# Cbam module
def cbam(x):
    x = channel_attention_module(x)
    x = spatial_attention_module(x)

    return x


def build_image_model():
    base_model = InceptionV3(
        weights='imagenet',
        include_top=False,
        input_shape=(299,299,3)
    )
    for layer in base_model.layers:
        layer.trainable = False
    x = cbam(base_model.output)

    x= GlobalAveragePooling2D()(x)
    predictions = Dense(1, activation='sigmoid')(x)

    model = Model(inputs=base_model.input, outputs=predictions)

    return model