# -- Code Cell --
import io
import ast
import cv2
import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt
from skimage.metrics import structural_similarity as ssim

# -- Code Cell --
# ==========================================
# TODO: PARTICIPANT LOGIC GOES HERE
# ==========================================

def compute_mean_difference(input_df): 
    # Task 1 placeholder
    df = input_df.copy()
    return 0

def reconstruct_image(img): 
    # Currently just returns the noisy image as-is
    return img
 
def make_prediction(input_df): 
    # Task 2 placeholder (Classification)
    df = input_df.copy()
    return ['algolymp'] * len(df)

# -- Code Cell --
import cv2 as cv

# -- Code Cell --
z = 8

# -- Code Cell --
fig = plt.figure(figsize=(8, 4), layout='tight')
ax1, ax2 = fig.subplots(1, 2)
noisy_img = cv.imread(f'data/images/train_data/image_{z}.jpg',cv.IMREAD_GRAYSCALE)
noisy_img_color = cv.imread(f'data/images/train_data/image_{z}.jpg',cv.IMREAD_COLOR)
target_img = cv.imread(f'data/images/validation_data/image_{z}.jpg',cv.IMREAD_COLOR)
noisy_img_color = cv.cvtColor(noisy_img_color, cv.COLOR_BGR2RGB)
target_img = cv.cvtColor(target_img, cv.COLOR_BGR2RGB)

b, g, r = cv2.split(noisy_img_color)
def shift_channel(channel, shift):
    h, w = channel.shape
    M = np.float32([[1, 0, shift], [0, 1, 0]])
    return cv2.warpAffine(channel, M, (w, h))
b_fixed = shift_channel(b, -5) 
r_fixed = shift_channel(r, 5)   
fixed = cv2.merge([b_fixed, g, r_fixed])

fixed = fixed[:, 5:-5]
noisy_img = noisy_img[:, 5:-5]
target_img = target_img[:, 5:-5]

#linia
t,reconstructed = cv.threshold(noisy_img,1,255,cv.THRESH_BINARY_INV)
mask = reconstructed

#blur
blur = cv.GaussianBlur(fixed,(7,7),0)

#vhs
img = np.zeros((noisy_img.shape[0],noisy_img.shape[1]), np.uint8)
for r in range(noisy_img.shape[0]):
        if r%2 != 0:
            cv.line(img,(0,r),(256,r),(255,255,255),1)
mask2 = img

#combinat
mask3 = cv.bitwise_or(mask, mask2)


#bright
bright = cv.convertScaleAbs(blur, alpha=1.0, beta=30)

#umplu
dstpt2 = cv.inpaint(bright,mask3,1,cv.INPAINT_TELEA)     
print(dstpt2.shape)



ax1.imshow(noisy_img_color,cmap="gray")
ax1.set_title("Clean / Target")
ax1.axis('off')

ax2.imshow(dstpt2,cmap="gray")
ax2.set_title("Reconstructed")
ax2.axis('off')

plt.show()

# Calculate Structural Similarity Index
score = compare_and_score(target_img, dstpt2)
print(f'Similarity Score: {score:.4f}')
print('------------------/------------------')

# -- Code Cell --
def compare_and_score(clean, reconstructed):
    """
    Visualizes side-by-side comparison and calculates SSIM score.
    """
    fig = plt.figure(figsize=(8, 4), layout='tight')
    ax1, ax2 = fig.subplots(1, 2)
    
    ax1.imshow(clean)
    ax1.set_title("Clean / Target")
    ax1.axis('off')
    
    ax2.imshow(reconstructed)
    ax2.set_title("Reconstructed")
    ax2.axis('off')
    
    plt.show()

    # Calculate Structural Similarity Index
    score = ssim(np.array(clean), np.array(reconstructed), channel_axis=2)
    print(f'Similarity Score: {score:.4f}')
    print('------------------/------------------')
    return score

# --- Local Validation Loop ---
validation_scores = []
for i in range(20):
    noisy_img = Image.open(f'data/images/train_data/image_{i}.jpg')
    target_img = Image.open(f'data/images/validation_data/image_{i}.jpg')

    # Run your custom reconstruction function
    reconstructed = reconstruct_image(noisy_img)
    
    # Check results
    score = compare_and_score(target_img, reconstructed)
    validation_scores.append(score)

print(f'Mean Validation Score: {np.mean(validation_scores):.4f}')

# -- Code Cell --
def generate_submission_file():
    """
    Compiles predictions for all subtasks into the required CSV format.
    """
    submission_rows = []
    test_df = pd.read_csv('data/test_data.csv')
 
    print("Processing Subtask 1: Mean Difference...")
    # (1, 1, Answer) format
    task1_result = compute_mean_difference(test_df)
    submission_rows.append((1, 1, task1_result))
    
    print("Processing Subtask 2: Classification...")
    preds = make_prediction(test_df)
    for idx, row in test_df.iterrows():
        # (2, ID, Label) format
        submission_rows.append((2, row['ID'], preds[idx])) 
 
    # Create and save DF
    submission_df = pd.DataFrame(submission_rows, columns=['subtaskID', 'datapointID', 'answer'])
    submission_df.to_csv('submission.csv', index=False)
    
    print("submission.csv created successfully.")
    return submission_df
 
submission = generate_submission_file()

# -- Code Cell --
