import glob

import numpy as np
import pandas as pd
from PIL import Image
from sklearn.model_selection import train_test_split


def create_pixel_csv(img_dir: str, save_dir: str, img_type: str):
    """
    Creates a CSV file with the pixel coordinates of the image
    :param img_dir: image directory containing the images
    :param save_dir: directory for saving the CSV file per image
    :param img_type: color or normal
    :return: saves the CSV file per image0
    """
    imgs = sorted(glob.glob(f"{img_dir}/*.png"))
    print(f"Found {len(imgs)} {img_type} images")
    for idx, img_path in enumerate(imgs):
        img = Image.open(img_path)
        img_np = np.array(img)
        xy_coords = np.flip(
            np.column_stack(np.where(np.all(img_np >= 0, axis=2))), axis=1
        )
        rgb = np.reshape(img_np, (np.prod(img_np.shape[:2]), 3))
        # Add pixel numbers in front
        pixel_numbers = np.expand_dims(np.arange(1, xy_coords.shape[0] + 1), axis=1)
        value = np.hstack([pixel_numbers, xy_coords, rgb])
        # Properly save as CSV
        if img_type == "color":
            df = pd.DataFrame(value, columns=["pixel_number", "X", "Y", "R", "G", "B"])
            del df["pixel_number"]
            df.to_csv(f"{save_dir}/image_{idx}.csv", sep=",", index=False)
        elif img_type == "normal":
            df = pd.DataFrame(value, columns=["pixel_number", "X", "Y", "Nx", "Ny", "Nz"])
            del df["pixel_number"]
            df.to_csv(f"{save_dir}/image_{idx}.csv", sep=",", index=False)
        print(f"{img_type} image CSV written for image {idx}")


def combine_csv(csv_dir: str, img_type: str):
    """
    Combines all the CSV files in the directory into one CSV file for later use in training

    :param csv_dir: directory containing the CSV files
    :param img_type: color or normal
    """
    csv_files=(glob.glob(f"{csv_dir}/*.csv"))
    print(f"Found {len(csv_files)} {img_type} CSV files")
    df = pd.concat([pd.read_csv(f, sep=",") for f in (glob.glob(f"{csv_dir}/*.csv"))])
    df.to_csv(f"{csv_dir}/combined.csv", sep=",", index=False)
    print(f"Combined CSV written for {img_type} images")
    check_nans(f"{csv_dir}/combined.csv")
    print("----------------------------------------------------")


def create_train_test_csv(save_dir: str, normal_path: str, color_path: str):
    """
    Creates a CSV file with the pixel coordinates of the image. Samples 4% from zeros.
    Splits the cleaned dataset into train and test (80%/20%)
    :param save_dir: path for train_test_split folder
    :param normal_path: path to combined normal CSV file
    :param color_path: path to combined color CSV file
    """
    seed=42
    rgb_df = pd.read_csv(color_path, sep=",")
    normal_df = pd.read_csv(normal_path, sep=",")

    rgb_df['Nx'] = normal_df['Nx']
    rgb_df['Ny'] = normal_df['Ny']
    rgb_df['Nz'] = normal_df['Nz']

    zeros_df = rgb_df[(rgb_df['Nx'] == 127) & (rgb_df['Ny'] == 127) ^ (rgb_df['Nz'] == 127)]
    print(f"Found {len(zeros_df)} zeros in the dataset")
    non_zeros_df = rgb_df[(rgb_df['Nx'] != 127) | (rgb_df['Ny'] != 127) & (rgb_df['Nz'] != 127)]
    print(f"Found {len(non_zeros_df)} non-zeros in the dataset")
    zeros_df.to_csv(f"{save_dir}/zeros.csv", sep=",", index=False)
    non_zeros_df.to_csv(f"{save_dir}/non_zeros.csv", sep=",", index=False)
    print("----------------------------------------------------")
    # sampling 4% of zeros
    clean_df = zeros_df.sample(frac=0.04, random_state=seed)
    clean_df = pd.concat([clean_df, non_zeros_df])
    clean_df.to_csv(f"{save_dir}/clean-data.csv", sep=",", index=False)
    print(f"Clean data CSV of {len(clean_df)} written")

    # split into train and test
    train_df,test_df=train_test_split(clean_df,test_size=0.2,random_state=seed)
    train_df.to_csv(f"{save_dir}/train.csv", sep=",", index=False)
    test_df.to_csv(f"{save_dir}/test.csv", sep=",", index=False)
    print(f"Train set of size {len(train_df)} and Test set of size {len(test_df)} written")


def check_nans(csv_path: str):
    """
    Checks if there are any NaNs in the CSV file
    :param csv_path: path to the CSV file
    :return: True if there are no NaNs, False otherwise
    """
    df = pd.read_csv(csv_path, sep=",")
    if df.isnull().values.any():
        nans=df.isnull().sum().sum()
        print(f"Found {nans} NaNs in {csv_path}")
        print("Replacing NaNs with mean")
        df.fillna(df.mean(), inplace=True)
        df.to_csv(csv_path, sep=",", index=False)
        print(f"NaNs replaced in {csv_path}")
    else:
        print("No NaNs found.Perfect!")