import cv2
import pathlib
import os

import numpy as np

from sample.preprocess import sliceAndMeanShift 

import matplotlib.pyplot as plt

'''
Database file structure example:

|-- ImageDatabase
│   |-- Images
|   |   |--<autoInc>-<imgExt>-<name> 
|   |   |   |--<autoInc>-<imgExt>-<name>-Original.<imgExt>
|   |   |   |--Sliced
|   |   |   |   |--128x128
|   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix
|   |   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix-Slice.<imgExt>  
|   |   |   |   |   |   |--Processed
|   |   |   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix-<procType>-Processed.<imgExt>
|   |   |   |   |--256x256
|   |   |   |   |--512x512
'''
def initDB(dbLoc: str, name: str) -> bool:
    '''
    Initialized the image database at a given directory
    File Struture:
    |-- ImageDatabase
    │   |-- Data
    '''
    dbLoc = pathlib.Path(dbLoc)
    dbDir = dbLoc.joinpath(name)
    if dbDir.is_dir():
        return False 
    else:
        imagesDir = dbDir.joinpath("Images")
        os.mkdir(dbDir)
        os.mkdir(imagesDir)
        return True



def importAndProcess(imgPath: str, dbPath: str, newImgName: str):
    '''
    Slices and segments an image and imports it into the database
    '''
    imgPath = pathlib.Path(imgPath)
    dbPath  = pathlib.Path(dbPath)

    # Ensure database exists
    if not dbPath.is_dir():
        print("ERROR: No database exists at {}".format(dbPath))
        return False

    # Internal database directory
    imagesPath = dbPath.joinpath("Images")

    # Ensure database was initialized correctly
    if not imagesPath.is_dir():
        print("ERROR - Database Not Initialized: Use 'initDB.py' to init database.")
        return False

    # PROCESS IMAGE

    # Step 1: Read in image
    imgArr  = cv2.imread(str(imgPath))
    fileExt = imgPath.name.split(".")[-1]

    # Auto increment id
    dbIDs = [int(f.name.split("-")[0]) for f in imagesPath.iterdir() ]
    dbID = 0  

    if len(dbIDs) != 0:
        dbID = max(dbIDs) + 1 

    # Create dir for image
    baseImagePath = imagesPath.joinpath("{}-{}-{}-OR".format(dbID, fileExt, newImgName))
    baseImageName = baseImagePath.name
    os.mkdir(baseImagePath)

    # Step 2: Save Original File
    baseImageFilePath = baseImagePath.joinpath("{}.{}".format(baseImageName, fileExt))
    cv2.imwrite(str(baseImageFilePath), imgArr)

    # Create dir for sliced images and models
    slicedImagePath = baseImagePath.joinpath("SL")
    os.mkdir(slicedImagePath)
    modelsPath = baseImagePath.joinpath("Models")
    os.mkdir(modelsPath)
    
    # Save Slices
    sizes = [ "128", "256", "512" ]

    for size in sizes:
        print("Processing size:", size)
        # Run processing
        processedResult = sliceAndMeanShift(imgArr, int(size))

        slicedSizePath = slicedImagePath.joinpath(size)
        os.mkdir(slicedSizePath)

        # Slices and corresponding processed images
        sliceList     = processedResult["slices"][size]
        processedList = processedResult["meanShift"][size]

        for i in range(len(sliceList)):
            # <autoInc>-<imgExt>-<name>-128-xPix-yPix
            x = sliceList[i][0]
            y = sliceList[i][1]
            sliceArr = sliceList[i][2]

            if sliceArr.max() != 0:
                sliceBaseName = "{}-{}-{}-{}-SL".format(baseImageName, size, x, y)
                sliceBasePath = slicedSizePath.joinpath(sliceBaseName)
                os.mkdir(sliceBasePath)

                # Save Sliced Image
                sliceFilePath = sliceBasePath.joinpath("{}.{}".format(sliceBaseName, fileExt))
                cv2.imwrite(str(sliceFilePath), sliceArr)

                # Create Processed Slice Path
                sliceProcPath = sliceBasePath.joinpath("PROC")
                os.mkdir(sliceProcPath)

                # Segmented image and label image
                msSegArr = processedList[i][3]
                msLabImg = processedList[i][4]

                msSegDir = sliceProcPath.joinpath("{}-{}-{}".format(sliceBaseName, 
                                                                    "MSSEG", 
                                                                    "PROC"))
                msSegFileDir = msSegDir.joinpath("{}.{}".format(msSegDir.name, fileExt))

                msLabFileDir = msSegDir.joinpath("{}-{}-{}".format(sliceBaseName, 
                                                                    "MSLAB", 
                                                                    "PROC"))

                os.mkdir(msSegDir)

                # Save processed image and label array
                cv2.imwrite(str(msSegFileDir), msSegArr) 
                np.save(str(msLabFileDir), msLabImg)