from samhi.data_processing.data_fetching import DataFetcher
from samhi.data_processing.data_utils import DataProcessing
from samhi.data_processing.dataset import HistologyDataset
from PIL import Image
import numpy as np
import tifffile
import torch
from tqdm import tqdm
import argparse
# datasets = ["CellSeg", "CoNIC", "CPM", "CRAG", "CryoNuSeg", "GlaS", 
#             "Janowczyk", "Kumar", "MoNuSAC", "NuClick", "PAIP2023", "SegPC"]

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--datasets", nargs='+', default=["CellSeg"])
    args = parser.parse_args()
    datasets = args.datasets

    data_directory = "/home/ubuntu/thesis/data"
    dataset = HistologyDataset(datasets, data_directory, "helmholtz", 1024, 1, "NoOp", 1)
    
    mask_count = 0
    for i in tqdm(range(1546)):
        gt_name = dataset[i][1][3][1]
        if gt_name.endswith(".tiff") or gt_name.endswith(".tif"):
            gt = tifffile.imread(gt_name)
        else:
            gt = Image.open(gt_name)
        gt = np.array(gt)
        gt_cca = DataProcessing.connected_component_analysis(gt, 1)
        mask_count += torch.max(gt_cca).item()
    print(datasets[0])
    print(mask_count)

if __name__ == "__main__":
    main()