import numpy as np
import cv2
import glob
from pathlib import Path
from matplotlib import pyplot as plt

from rescale import *
from slid import detect_lines
from laps import LAPS
from llr import LLR, llr_pad

RAW_DATA_FOLDER = './data/raw/games/'
PREPROCESSED_FOLDER = './data/preprocessed/games/'


def preprocess_image(path, final_folder="", filename="",  save=False):
    res = cv2.imread(path)[..., ::-1]
    for _ in range(2):
        img, shape, scale = image_resize(res)
        lines = detect_lines(img)
        lattice_points = LAPS(img, lines)
        inner_points = LLR(img, lattice_points, lines)
        four_points = llr_pad(inner_points, img)
        try:
            res = crop(res, four_points, scale)
        except:
            print("WARNING: couldn't crop around outer points")
            res = crop(
                res, inner_points, scale)
    if save:
        Path(final_folder).mkdir(parents=True, exist_ok=True)
        plt.imsave("%s/%s" % (final_folder, filename), res)
    return res


def preprocess_single_image(game_name, version, image_number):
    raw_path = RAW_DATA_FOLDER + '%s/%s/%i.*' % (game_name, version, image_number)
    final_folder = PREPROCESSED_FOLDER + "%s/%s/" % (game_name, version)
    
    matching_files = glob.glob(raw_path)
    
    if not matching_files:
        print(f"ERROR: No image found at {raw_path}")
        return False
    
    if len(matching_files) > 1:
        print(f"WARNING: Multiple files found, using first one: {matching_files[0]}")
    
    path = matching_files[0]
    print(f"Processing: {path}")
    
    try:
        preprocess_image(path, final_folder=final_folder,
                        filename="%i.png" % image_number, save=True)
        print(f"✓ Successfully saved to {final_folder}{image_number}.png")
        return True
    except Exception as e:
        print(f"✗ ERROR processing image: {e}")
        import traceback
        traceback.print_exc()
        return False


def preprocess_multiple_images(images_list):
    total = len(images_list)
    success = 0
    failed = []
    
    print(f"Processing {total} images...")
    for game_name, version, img_num in images_list:
        print(f"\n[{success + len(failed) + 1}/{total}] {game_name}/{version}/image {img_num}")
        if preprocess_single_image(game_name, version, img_num):
            success += 1
        else:
            failed.append((game_name, version, img_num))
    
    print(f"\n{'='*50}")
    print(f"Summary: {success}/{total} images processed successfully")
    if failed:
        print(f"Failed images ({len(failed)}):")
        for game, ver, num in failed:
            print(f"  - {game}/{ver}/image {num}")


def preprocess_games(game_list):
    for game_name in game_list:
        for ver in ['orig', 'rev']:
            img_filename_list = []
            folder_name = RAW_DATA_FOLDER + '%s/%s/*' % (game_name, ver)
            for path_name in glob.glob(folder_name):
                img_filename_list.append(path_name)

            final_folder = PREPROCESSED_FOLDER + "%s/%s/" % (game_name, ver)
            count = 0
            img_filename_list.sort(key=lambda s: int(
                s.split('/')[-1].split('.')[0]))
            for path in img_filename_list:
                count += 1
                preprocess_image(path, final_folder=final_folder,
                                 filename="%i.png" % count, save=True)
            if count > 0:
                print("Done saving in %s." % final_folder)
            else:
                print("No images found in %s" % folder_name)


if __name__ == '__main__':
    # deja fait: carlsen_anand_2014, carlsen_gukesh_2025 vachier-lagrave_carlsen_2023 david_vachier-lagrave_2014
    # game_list = ['david_vachier-lagrave_2014']
    # preprocess_games(game_list)
    preprocess_single_image('vachier-lagrave_carlsen_2023', 'rev', 57)