import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
import pywt  # Thư viện xử lý wavelet
from PIL import Image
from tensorflow.keras import layers, models
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
import joblib
import io
import os
import cv2
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers, models
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import matplotlib.pyplot as plt
import random
from keras.applications import ResNet50
from tensorflow.keras.applications.resnet import preprocess_input
from tensorflow.keras import layers, models
from tensorflow.keras.applications.resnet import preprocess_input
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras.layers import Lambda  # Đảm bảo nhập Lambda từ tensorflow.keras.layers
from skimage.feature import graycomatrix, graycoprops
from keras.applications import ResNet50
from tensorflow.keras.applications.resnet import preprocess_input

# Add Cloudinary import
import cloudinary
import cloudinary.uploader
from cloudinary.utils import cloudinary_url

# Cloudinary Configuration
cloudinary.config(
    cloud_name = os.getenv("CLOUD"),
    api_key = os.getenv("API"),
    api_secret = os.getenv("SECRET"),
    secure=True
)

def upload_to_cloudinary(file_path, label):
    """
    Upload file to Cloudinary with specified label as folder
    """
    try:
        # Upload to Cloudinary
        upload_result = cloudinary.uploader.upload(
            file_path,
            folder=label,
            public_id=f"{label}_{os.path.basename(file_path)}"
        )

        # Generate optimized URLs
        optimize_url, _ = cloudinary_url(
            upload_result['public_id'],
            fetch_format="auto",
            quality="auto"
        )

        auto_crop_url, _ = cloudinary_url(
            upload_result['public_id'],
            width=500,
            height=500,
            crop="auto",
            gravity="auto"
        )

        return {
            "upload_result": upload_result,
            "optimize_url": optimize_url,
            "auto_crop_url": auto_crop_url
        }

    except Exception as e:
        return f"Error uploading to Cloudinary: {str(e)}"

def main():
    st.title("Web App Phân loại đá")
    st.write("Tải lên hình ảnh của một viên đá để phân loại loại của nó.")

    # Load model and scaler
    model, scaler = load_model_and_scaler()
    if model is None or scaler is None:
        st.error("Không thể tải mô hình hoặc bộ chuẩn hóa. Vui lòng đảm bảo rằng cả hai tệp đều tồn tại.")
        return

    # Initialize session state
    if 'predictions' not in st.session_state:
        st.session_state.predictions = None
    if 'uploaded_image' not in st.session_state:
        st.session_state.uploaded_image = None

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Tải lên Hình ảnh")
        uploaded_file = st.file_uploader("Chọn hình ảnh...", type=["jpg", "jpeg", "png"])

        if uploaded_file is not None:
            try:
                image = Image.open(uploaded_file)
                st.image(image, caption="Hình ảnh đã tải lên", use_column_width=True)
                st.session_state.uploaded_image = image

                with st.spinner('Đang phân tích hình ảnh...'):
                    processed_image = preprocess_image(image, scaler)
                    prediction = model.predict(processed_image)

                    class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
                    st.session_state.predictions = get_top_predictions(prediction, class_names)

            except Exception as e:
                st.error(f"Lỗi khi xử lý hình ảnh: {str(e)}")

    with col2:
        st.subheader("Kết quả Dự đoán")
        if st.session_state.predictions:
            # Display main prediction
            top_class, top_confidence = st.session_state.predictions[0]
            st.markdown(
                f"""
                <div class='prediction-card'>
                    <h3>Dự đoán chính: Màu {top_class}</h3>
                    <h3>Độ tin cậy: {top_confidence:.2f}%</h3>
                </div>
                """,
                unsafe_allow_html=True
            )

            # Display confidence bar
            st.progress(float(top_confidence) / 100)

            # Display top 5 predictions
            st.markdown("### 5 Dự đoán hàng đầu")
            st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)

            for class_name, confidence in st.session_state.predictions:
                st.markdown(
                    f"**Màu {class_name}: Độ tin cậy {confidence:.2f}%**"
                )
                st.progress(float(confidence) / 100)

            st.markdown("</div>", unsafe_allow_html=True)

            # User Confirmation Section
            st.markdown("### Xác nhận độ chính xác của mô hình")
            st.write("Giúp chúng tôi cải thiện mô hình bằng cách xác nhận độ chính xác của dự đoán.")

            # Accuracy Radio Button
            accuracy_option = st.radio(
                "Dự đoán có chính xác không?",
                ["Chọn", "Chính xác", "Không chính xác"],
                index=0
            )

            if accuracy_option == "Không chính xác":
                # Input for correct grade
                correct_grade = st.selectbox(
                    "Chọn màu đá đúng:",
                    ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7'],
                    index=None,
                    placeholder="Chọn màu đúng"
                )

                # Upload button
                if st.button("Tải lên Hình ảnh để sửa chữa"):
                    if correct_grade and st.session_state.uploaded_image:
                        # Save the image temporarily
                        temp_image_path = f"temp_image_{hash(uploaded_file.name)}.png"
                        st.session_state.uploaded_image.save(temp_image_path)

                        try:
                            # Upload to Cloudinary
                            cloudinary_result = upload_to_cloudinary(temp_image_path, correct_grade)

                            if isinstance(cloudinary_result, dict):
                                st.success(f"Hình ảnh đã được tải lên thành công cho màu {correct_grade}")
                                st.write(f"URL công khai: {cloudinary_result['upload_result']['secure_url']}")
                            else:
                                st.error(cloudinary_result)

                            # Clean up temporary file
                            os.remove(temp_image_path)

                        except Exception as e:
                            st.error(f"Tải lên thất bại: {str(e)}")
                    else:
                        st.warning("Vui lòng chọn màu đúng trước khi tải lên.")
        else:
            st.info("Tải lên hình ảnh để xem các dự đoán.")

    st.markdown("---")
    st.markdown("Tạo bởi ❤️ với Streamlit")

def load_model_and_scaler():
    """Load the trained model and scaler"""
    try:
        model = tf.keras.models.load_model('mlp_model.h5')
        # Tải scaler đã lưu
        scaler = joblib.load('scaler.pkl')
        return model, scaler
    except Exception as e:
        st.error(f"Error loading model or scaler: {str(e)}")
        return None, None

def color_histogram(image, bins=16):
    """
    Tính histogram màu cho ảnh RGB

    Args:
        image (np.ndarray): Ảnh đầu vào
        bins (int): Số lượng bins của histogram

    Returns:
        np.ndarray: Histogram màu được chuẩn hóa
    """
    # Kiểm tra và chuyển đổi ảnh
    if image is None or image.size == 0:
        raise ValueError("Ảnh không hợp lệ")

    # Đảm bảo ảnh ở dạng uint8
    if image.dtype != np.uint8:
        image = (image * 255).astype(np.uint8)

    # Tính histogram cho từng kênh màu
    hist_r = cv2.calcHist([image], [0], None, [bins], [0, 256]).flatten()
    hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten()
    hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten()

    # Chuẩn hóa histogram
    hist_r = hist_r / np.sum(hist_r) if np.sum(hist_r) > 0 else hist_r
    hist_g = hist_g / np.sum(hist_g) if np.sum(hist_g) > 0 else hist_g
    hist_b = hist_b / np.sum(hist_b) if np.sum(hist_b) > 0 else hist_b

    return np.concatenate([hist_r, hist_g, hist_b])

def color_moments(image):
    """
    Tính các moment màu cho ảnh

    Args:
        image (np.ndarray): Ảnh đầu vào

    Returns:
        np.ndarray: Các moment màu
    """
    # Kiểm tra và chuyển đổi ảnh
    if image is None or image.size == 0:
        raise ValueError("Ảnh không hợp lệ")

    # Đảm bảo ảnh ở dạng float và chuẩn hóa
    img = image.astype(np.float32) / 255.0 if image.max() > 1 else image.astype(np.float32)

    moments = []
    for i in range(3):  # Cho mỗi kênh màu
        channel = img[:,:,i]

        # Tính các moment
        mean = np.mean(channel)
        std = np.std(channel)
        skewness = np.mean(((channel - mean) / (std + 1e-8)) ** 3)

        moments.extend([mean, std, skewness])

    return np.array(moments)

def dominant_color_descriptor(image, k=3):
    """
    Xác định các màu chính thống trị trong ảnh

    Args:
        image (np.ndarray): Ảnh đầu vào
        k (int): Số lượng màu chủ đạo

    Returns:
        np.ndarray: Các màu chủ đạo và tỷ lệ
    """
    # Kiểm tra và chuyển đổi ảnh
    if image is None or image.size == 0:
        raise ValueError("Ảnh không hợp lệ")

    # Đảm bảo ảnh ở dạng uint8
    if image.dtype != np.uint8:
        image = (image * 255).astype(np.uint8)

    # Reshape ảnh thành mảng pixel
    pixels = image.reshape(-1, 3)

    # Các tham số cho K-means
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    flags = cv2.KMEANS_RANDOM_CENTERS

    try:
        # Thực hiện phân cụm K-means
        _, labels, centers = cv2.kmeans(
            pixels.astype(np.float32), k, None, criteria, 10, flags
        )

        # Tính toán số lượng và tỷ lệ của từng cụm
        unique, counts = np.unique(labels, return_counts=True)
        percentages = counts / len(labels)

        # Kết hợp các màu và tỷ lệ
        dominant_colors = centers.flatten()
        color_percentages = percentages

        return np.concatenate([dominant_colors, color_percentages])
    except Exception:
        # Trả về mảng 0 nếu có lỗi
        return np.zeros(2 * k)

def color_coherence_vector(image, k=3):
    """
    Tính vector liên kết màu

    Args:
        image (np.ndarray): Ảnh đầu vào
        k (int): Số lượng vùng

    Returns:
        np.ndarray: Vector liên kết màu
    """
    # Kiểm tra và chuyển đổi ảnh
    if image is None or image.size == 0:
        raise ValueError("Ảnh không hợp lệ")

    # Chuyển sang ảnh xám
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Đảm bảo ảnh ở dạng uint8
    if gray.dtype != np.uint8:
        gray = np.uint8(gray)

    # Áp dụng Otsu's thresholding
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # Phân tích thành phần liên thông
    num_labels, labels = cv2.connectedComponents(binary)

    ccv = []
    for i in range(1, min(k+1, num_labels)):
        region_mask = (labels == i)
        total_pixels = np.sum(region_mask)
        coherent_pixels = total_pixels

        ccv.extend([coherent_pixels, total_pixels])

    # Đảm bảo độ dài vector
    while len(ccv) < 2 * k:
        ccv.append(0)

    return np.array(ccv)

def edge_features(image, bins=16):
    """
    Trích xuất đặc trưng cạnh từ ảnh

    Args:
        image (np.ndarray): Ảnh đầu vào
        bins (int): Số lượng bins của histogram

    Returns:
        np.ndarray: Đặc trưng cạnh
    """
    # Kiểm tra và chuyển đổi ảnh
    if image is None or image.size == 0:
        raise ValueError("Ảnh không hợp lệ")

    # Chuyển sang ảnh xám
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Đảm bảo ảnh ở dạng uint8
    if gray.dtype != np.uint8:
        gray = np.uint8(gray)

    # Tính Sobel edges
    sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    sobel_mag = np.sqrt(sobel_x**2 + sobel_y**2)

    # Chuẩn hóa độ lớn Sobel
    sobel_mag = np.uint8(255 * sobel_mag / np.max(sobel_mag))

    # Tính histogram của Sobel magnitude
    sobel_hist = cv2.calcHist([sobel_mag], [0], None, [bins], [0, 256]).flatten()
    sobel_hist = sobel_hist / np.sum(sobel_hist) if np.sum(sobel_hist) > 0 else sobel_hist

    # Tính mật độ cạnh bằng Canny
    canny_edges = cv2.Canny(gray, 100, 200)
    edge_density = np.sum(canny_edges) / (gray.shape[0] * gray.shape[1])

    return np.concatenate([sobel_hist, [edge_density]])



def histogram_in_color_space(image, color_space='HSV', bins=16):
    """
    Tính histogram của ảnh trong một không gian màu mới.
    """
    if color_space == 'HSV':
        converted = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    elif color_space == 'LAB':
        converted = cv2.cvtColor(image, cv2.COLOR_RGB2Lab)
    else:
        raise ValueError("Unsupported color space")

    histograms = []
    for i in range(3):  # 3 kênh màu
        hist = cv2.calcHist([converted], [i], None, [bins], [0, 256]).flatten()
        hist = hist / np.sum(hist)
        histograms.append(hist)

    return np.concatenate(histograms)

def glcm_features(image, distances=[1, 2, 3], angles=[0, np.pi/4, np.pi/2, 3*np.pi/4], levels=256):
    """
    Tính các đặc trưng GLCM của ảnh grayscale.
    """
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Đảm bảo ảnh ở dạng uint8
    if gray.dtype != np.uint8:
        gray = (gray * 255).astype(np.uint8)

    glcm = graycomatrix(gray, distances=distances, angles=angles, levels=levels, symmetric=True, normed=True)

    features = []
    # Các thuộc tính phổ biến: contrast, homogeneity, energy, correlation
    for prop in ['contrast', 'homogeneity', 'energy', 'correlation']:
        features.extend(graycoprops(glcm, prop).flatten())

    return np.array(features)


def gabor_features(image, kernels=None):
    """
    Tính các đặc trưng từ bộ lọc Gabor.
    """
    if kernels is None:
        kernels = []
        for theta in np.arange(0, np.pi, np.pi / 4):  # Các góc từ 0 đến 180 độ
            for sigma in [1, 3]:  # Các giá trị sigma
                for frequency in [0.1, 0.5]:  # Các tần số
                    kernel = cv2.getGaborKernel((9, 9), sigma, theta, 1/frequency, gamma=0.5, ktype=cv2.CV_32F)
                    kernels.append(kernel)

    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    features = []
    for kernel in kernels:
        filtered = cv2.filter2D(gray, cv2.CV_32F, kernel)
        features.append(filtered.mean())
        features.append(filtered.var())

    return np.array(features)

def wavelet_features(image, wavelet='db1', level=3):
    """
    Trích xuất các hệ số wavelet từ ảnh grayscale.
    """
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    coeffs = pywt.wavedec2(gray, wavelet, level=level)
    features = []
    for coeff in coeffs:
        if isinstance(coeff, tuple):  # Chi tiết (LH, HL, HH)
            for subband in coeff:
                features.append(subband.mean())
                features.append(subband.var())
        else:  # Xấp xỉ (LL)
            features.append(coeff.mean())
            features.append(coeff.var())

    return np.array(features)

def fractal_dimension(image):
    """
    Tính Fractal Dimension của ảnh.
    """
    # Chuyển đổi ảnh sang grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Đảm bảo ảnh ở dạng uint8
    if gray.dtype != np.uint8:
        gray = (gray * 255).astype(np.uint8)

    # Áp dụng Canny để tìm cạnh
    edges = cv2.Canny(gray, 100, 200)

    # Tính fractal dimension dựa trên phương pháp box-counting
    sizes = []
    counts = []
    for size in range(2, 65, 2):  # Kích thước hộp từ 2 đến 64
        region_size = (edges.shape[0] // size, edges.shape[1] // size)
        count = np.sum(cv2.resize(edges, region_size, interpolation=cv2.INTER_AREA) > 0)
        sizes.append(size)
        counts.append(count)

    # Tính log-log slope
    log_sizes = np.log(sizes)
    log_counts = np.log(counts)
    slope, _ = np.polyfit(log_sizes, log_counts, 1)

    # Trả về giá trị fractal dimension
    return np.array([slope])


def extract_features(image):
    """
    Extract multiple features from an image, including edge-based features.
    """
    color_hist = color_histogram(image)
    color_mom = color_moments(image)
    dom_color = dominant_color_descriptor(image)
    ccv = color_coherence_vector(image)
    edges = edge_features(image)

    # Các đặc trưng từ phương pháp mới
    hsv_hist = histogram_in_color_space(image, color_space='HSV')
    # lab_hist = histogram_in_color_space(image, color_space='LAB')
    glcm = glcm_features(image)
    gabor = gabor_features(image)
    wavelet = wavelet_features(image)
    # fractal = fractal_dimension(image)

    # Kết hợp tất cả thành một vector đặc trưng
    return np.concatenate([
        color_hist,
        color_mom,
        dom_color,
        ccv,
        edges,
        hsv_hist,
        # lab_hist,
        glcm,
        gabor,
        wavelet,
        # fractal
    ])


def create_resnet50_feature_extractor(input_shape=(256, 256, 3), num_classes=None):
    # Xây dựng mô hình ResNet112 đã huấn luyện sẵn từ TensorFlow
    inputs = layers.Input(shape=input_shape)

    # Thêm lớp Lambda để tiền xử lý ảnh
    x = Lambda(preprocess_input, output_shape=input_shape)(inputs)  # Xử lý ảnh đầu vào

    # Sử dụng mô hình ResNet112 đã được huấn luyện sẵn
    resnet50_model = ResNet50(include_top=False, weights='imagenet', input_tensor=x)

    # Trích xuất đặc trưng từ mô hình ResNet112
    x = layers.GlobalAveragePooling2D()(resnet50_model.output)

    if num_classes:
        x = layers.Dense(num_classes, activation='softmax')(x)  # Thêm lớp phân loại (nếu có)

    return models.Model(inputs=inputs, outputs=x)

def extract_features(image):
    """
    Extract multiple features from an image, including edge-based features.
    """
    color_hist = color_histogram(image)
    color_mom = color_moments(image)
    dom_color = dominant_color_descriptor(image)
    ccv = color_coherence_vector(image)
    edges = edge_features(image)

    # Các đặc trưng từ phương pháp mới
    hsv_hist = histogram_in_color_space(image, color_space='HSV')
    # lab_hist = histogram_in_color_space(image, color_space='LAB')
    glcm = glcm_features(image)
    gabor = gabor_features(image)
    wavelet = wavelet_features(image)
    # fractal = fractal_dimension(image)

    # Kết hợp tất cả thành một vector đặc trưng
    return np.concatenate([
        color_hist,
        color_mom,
        dom_color,
        ccv,
        edges,
        hsv_hist,
        # lab_hist,
        glcm,
        gabor,
        wavelet,
        # fractal
    ])

def preprocess_image(image, scaler):
    image=np.array(image)
    img_size=(256, 256)
    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, img_size)
    img_array = img.astype('float32') / 255.0

    features1 = np.array(extract_features(img_array))
    resnet_extractor = create_resnet50_feature_extractor()
    features2 = resnet_extractor.predict(np.expand_dims(img_array, axis=0))
    print(f"Shape of features1: {features1.shape}")
    print(f"Shape of features2: {features2.shape}")
    features = np.concatenate([np.expand_dims(features1, axis=0), features2], axis=1)  # Concatenate along axis 0

    # Scale features using the provided scaler
    scaled_features = scaler.transform(features)  # Reshape for scaling

    return scaled_features

def get_top_predictions(prediction, class_names):
    # Extract the top 5 predictions with confidence values
    probabilities = tf.nn.softmax(prediction[0]).numpy()
    top_indices = np.argsort(probabilities)[-5:][::-1]
    return [(class_names[i], probabilities[i] * 100) for i in top_indices]

if __name__ == "__main__":
    main()