import streamlit as st
import numpy as np
import cv2
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import plotly.express as px

# Dummy CNN Model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x1 = F.relu(self.conv1(x))  # First conv layer activation
        x2 = F.relu(self.conv2(x1))
        x3 = F.adaptive_avg_pool2d(x2, (8, 8))
        x4 = x3.view(x3.size(0), -1)
        x5 = F.relu(self.fc1(x4))
        x6 = self.fc2(x5)
        return x6, x1  # Return both output and first layer activations

# FFT processing functions
def apply_fft(image):
    fft_channels = []
    for channel in cv2.split(image):
        fft = np.fft.fft2(channel)
        fft_shifted = np.fft.fftshift(fft)
        fft_channels.append(fft_shifted)
    return fft_channels

def filter_fft_percentage(fft_channels, percentage):
    filtered_fft = []
    for fft_data in fft_channels:
        magnitude = np.abs(fft_data)
        sorted_mag = np.sort(magnitude.flatten())[::-1]
        num_keep = int(len(sorted_mag) * percentage / 100)
        threshold = sorted_mag[num_keep - 1] if num_keep > 0 else 0
        mask = magnitude >= threshold
        filtered_fft.append(fft_data * mask)
    return filtered_fft

def inverse_fft(filtered_fft):
    reconstructed_channels = []
    for fft_data in filtered_fft:
        fft_ishift = np.fft.ifftshift(fft_data)
        img_reconstructed = np.fft.ifft2(fft_ishift).real
        img_normalized = cv2.normalize(img_reconstructed, None, 0, 255, cv2.NORM_MINMAX)
        reconstructed_channels.append(img_normalized.astype(np.uint8))
    return cv2.merge(reconstructed_channels)

# CNN Pass Visualization
def pass_to_cnn(fft_image):
    model = SimpleCNN()
    magnitude_tensor = torch.tensor(np.abs(fft_image), dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    
    with torch.no_grad():
        output, activations = model(magnitude_tensor)
    
    # Ensure activations have the correct shape [batch_size, channels, height, width]
    if len(activations.shape) == 3:
        activations = activations.unsqueeze(0)  # Add batch dimension if missing
    
    return activations, magnitude_tensor

# 3D plotting function
def create_3d_plot(fft_channels, downsample_factor=1):
    fig = make_subplots(
        rows=3, cols=2,
        specs=[[{'type': 'scene'}, {'type': 'scene'}],
               [{'type': 'scene'}, {'type': 'scene'}],
               [{'type': 'scene'}, {'type': 'scene'}]],
        subplot_titles=(
            'Blue - Magnitude', 'Blue - Phase',
            'Green - Magnitude', 'Green - Phase',
            'Red - Magnitude', 'Red - Phase'
        )
    )

    for i, fft_data in enumerate(fft_channels):
        fft_down = fft_data[::downsample_factor, ::downsample_factor]
        magnitude = np.abs(fft_down)
        phase = np.angle(fft_down)
        
        rows, cols = magnitude.shape
        x = np.linspace(-cols//2, cols//2, cols)
        y = np.linspace(-rows//2, rows//2, rows)
        X, Y = np.meshgrid(x, y)

        fig.add_trace(
            go.Surface(x=X, y=Y, z=magnitude, colorscale='Viridis', showscale=False),
            row=i+1, col=1
        )
        
        fig.add_trace(
            go.Surface(x=X, y=Y, z=phase, colorscale='Inferno', showscale=False),
            row=i+1, col=2
        )

    fig.update_layout(
        height=1500,
        width=1200,
        margin=dict(l=0, r=0, b=0, t=30),
        scene_camera=dict(eye=dict(x=1.5, y=1.5, z=0.5)),
        scene=dict(
            xaxis=dict(title='Frequency X'),
            yaxis=dict(title='Frequency Y'),
            zaxis=dict(title='Magnitude/Phase')
        )
    )
    return fig

# Streamlit UI
st.set_page_config(layout="wide")
st.title("Interactive Frequency Domain Analysis with CNN")

# Initialize session state
if 'fft_channels' not in st.session_state:
    st.session_state.fft_channels = None
if 'filtered_fft' not in st.session_state:
    st.session_state.filtered_fft = None
if 'reconstructed' not in st.session_state:
    st.session_state.reconstructed = None
if 'show_cnn' not in st.session_state:
    st.session_state.show_cnn = False

# Upload image
uploaded_file = st.file_uploader("Upload an image", type=['png', 'jpg', 'jpeg'])

if uploaded_file is not None:
    file_bytes = np.frombuffer(uploaded_file.getvalue(), np.uint8)
    image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    st.image(image_rgb, caption="Original Image", use_column_width=True)

    # Apply FFT and store in session state
    if st.session_state.fft_channels is None:
        st.session_state.fft_channels = apply_fft(image)

    # Frequency percentage slider
    percentage = st.slider(
        "Percentage of frequencies to retain:",
        0.1, 100.0, 10.0, 0.1,
        help="Adjust the slider to select what portion of frequency components to keep."
    )

    # Apply FFT filter
    if st.button("Apply Filter"):
        st.session_state.filtered_fft = filter_fft_percentage(st.session_state.fft_channels, percentage)
        st.session_state.reconstructed = inverse_fft(st.session_state.filtered_fft)
        st.session_state.show_cnn = False  # Reset CNN visualization

    # Display reconstructed image and FFT data
    if st.session_state.reconstructed is not None:
        reconstructed_rgb = cv2.cvtColor(st.session_state.reconstructed, cv2.COLOR_BGR2RGB)
        st.image(reconstructed_rgb, caption="Reconstructed Image", use_column_width=True)

        # FFT Data Tables
        st.subheader("Frequency Data of Each Channel")
        for i, channel_name in enumerate(['Blue', 'Green', 'Red']):
            st.write(f"### {channel_name} Channel FFT Data")
            magnitude_df = pd.DataFrame(np.abs(st.session_state.filtered_fft[i]))
            phase_df = pd.DataFrame(np.angle(st.session_state.filtered_fft[i]))
            st.write("#### Magnitude Data:")
            st.dataframe(magnitude_df.head(10))
            st.write("#### Phase Data:")
            st.dataframe(phase_df.head(10))

        # 3D Visualization
        st.subheader("3D Frequency Components Visualization")
        downsample = st.slider(
            "Downsampling factor for 3D plots:",
            1, 20, 5,
            help="Controls the resolution of the 3D surface plots."
        )
        fig = create_3d_plot(st.session_state.filtered_fft, downsample)
        st.plotly_chart(fig, use_container_width=True)

        # Custom CSS to style the button
        st.markdown("""
            <style>
                .centered-button {
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    margin-top: 20px;
                }
                .stButton>button {
                    padding: 20px 40px;
                    font-size: 20px;
                    background-color: #4CAF50;
                    color: white;
                    border: none;
                    border-radius: 10px;
                    cursor: pointer;
                }
                .stButton>button:hover {
                    background-color: #45a049;
                }
            </style>
        """, unsafe_allow_html=True)

        # CNN Visualization Section
        with st.container():
            st.markdown('<div class="centered-button">', unsafe_allow_html=True)
            if st.button("Pass to CNN"):
                st.session_state.show_cnn = True
            st.markdown('</div>', unsafe_allow_html=True)

        if st.session_state.show_cnn:
            st.subheader("CNN Processing Visualization")
            activations, magnitude_tensor = pass_to_cnn(st.session_state.filtered_fft[0])
            
            # Display input tensor
            st.write("### Input Magnitude Tensor:")
            st.image(magnitude_tensor.squeeze().numpy(), 
                    caption="Magnitude Tensor", 
                    use_column_width=True,
                    clamp=True)
            
            # Display activations with improved visualization
            st.write("### First Convolution Layer Activations")
            activation = activations.detach().numpy()
            
            if len(activation.shape) == 4:
                # Create a grid of activation maps
                cols = 4  # Number of columns in the grid
                rows = 4  # 16 channels / 4 columns = 4 rows
                fig, axs = plt.subplots(rows, cols, figsize=(20, 20))
                
                for i in range(activation.shape[1]):
                    act_img = activation[0, i, :, :]
                    ax = axs[i//cols, i%cols]
                    ax.imshow(act_img, cmap='viridis')
                    ax.set_title(f'Channel {i+1}')
                    ax.axis('off')
                
                st.pyplot(fig)
                
                # Display sample activation values
                st.write("### Activation Values Sample")
                sample_activation = activation[0, 0, :10, :10]  # First 10x10 values
                st.dataframe(pd.DataFrame(sample_activation))
            
            # Additional Steps After Activation Channels
            st.markdown("---")
            st.subheader("Next Processing Steps in CNN")
            
            # Step 2: Second Convolution Layer Visualization
            st.write("### Second Convolution Layer Features")
            with torch.no_grad():
                model = SimpleCNN()
                output, activations = model(magnitude_tensor)
                second_conv = model.conv2(activations).detach().numpy()
            
            if len(second_conv.shape) == 4:
                cols = 8  # 32 channels / 8 columns = 4 rows
                rows = 4
                fig2, axs2 = plt.subplots(rows, cols, figsize=(20, 10))
                
                for i in range(second_conv.shape[1]):
                    act_img = second_conv[0, i, :, :]
                    ax = axs2[i//cols, i%cols]
                    ax.imshow(act_img, cmap='plasma')
                    ax.set_title(f'Channel {i+1}')
                    ax.axis('off')
                
                st.pyplot(fig2)
            
            # Step 3: Pooling Layer Visualization
            st.write("### Adaptive Average Pooling Output")
            with torch.no_grad():
                pooled = F.adaptive_avg_pool2d(torch.tensor(second_conv), (8, 8)).numpy()

            st.write("Pooled Features Shape:", pooled.shape)

            # Normalize and display pooled features
            pooled_sample = pooled[0, 0]
            pooled_normalized = (pooled_sample - pooled_sample.min()) / (pooled_sample.max() - pooled_sample.min())
            st.image(pooled_normalized, 
                    caption="Sample Pooled Feature Map", 
                    use_container_width=True,
                    clamp=True)
            
            # Step 4: Final Classification
            st.write("### Final Classification Scores")
            with torch.no_grad():
                model = SimpleCNN()
                output, _ = model(magnitude_tensor)
                scores = F.softmax(output, dim=1).numpy()

            classes = [f"Class {i}" for i in range(10)]
            fig3 = px.bar(x=classes, y=scores[0], title="Classification Probabilities")
            st.plotly_chart(fig3)
            
            # Step 5: Full Process Explanation
            st.markdown("""
            #### Processing Pipeline:
            1. Input Magnitude Spectrum → 
            2. Conv1 Features (16 channels) → 
            3. Conv2 Features (32 channels) → 
            4. Pooled Features → 
            5. Fully Connected Layers → 
            6. Final Classification
            """)