import os
import glob
import yaml
import re
import argparse




#creation of yaml file 

def yaml_file(

        images_dir:str,
        class_names: list[str],
        yaml_path:str

):
    """
    Creates a YAML file with information about the image directories, number of classes, and class names.

    Args:
    - images_dir (str): Path to the root directory containing subdirectories with images.
    - class_names (list[str]): A list of class names.
    - yaml_path (str): Path to the YAML file that will be created.

    Returns:
    - None
    """

    train_dir=os.path.join(images_dir,"train")
    test_dir=os.path.join(images_dir,"test")
    val_dir=os.path.join(images_dir,"val")

    class_number = len(class_names)


    with open(yaml_path,"w") as f:
        yaml.dump({
            "train":train_dir,
            "test":test_dir,
            "val":val_dir,
            "nc": class_number,
            "names":class_names


        },f)


def read_class_names_from_file(class_name_file):

    assert os.path.isfile(class_name_file),f"class name file{class_name_file} not exist"

    with open(class_name_file,"r") as f:
        file_content=f.read()
    
    class_names=file_content.split("\n")

    return class_names

def get_max_class_number(labels_dir):
    max_class_number=0
    label_glob=glob.glob(os.path.join(labels_dir, "**","*.txt"))
    pattern =r"^(\d)+([0-9]+\.?[0-9]*){4}$"
    for label_file in label_glob:
        with open(label_file,"r") as f:
            for line in f:
                matches= re.match(pattern, line)
                if matches:
                    max_class_number=max(max_class_number,matches.group(1))


    return max_class_number+1


def create_yaml(model_dir,images_dir,labels_dir):
    model_dir=os.path.abspath(model_dir)
    images_dir=os.path.abspath(images_dir)
    labels_dir=os.path.abspath(labels_dir)
  
    
   
    

    class_name_file=os.path.join(model_dir,"classes.txt")

    if os.path.isfile(class_name_file):
        class_names=list(read_class_names_from_file(class_name_file))
    else:
        max_class_number=get_max_class_number(labels_dir=labels_dir)
        class_names=list(map(lambda x: "Default"+str(x) , range(max_class_number)))

    yaml_path=os.path.join(model_dir,"data.yaml")
    yaml_file(images_dir=images_dir,
              class_names=class_names,
              yaml_path=yaml_path)
    
    print("yaml file created")
    
    
    
    
if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Create YAML file for YOLOv5')
    parser.add_argument('--model_dir', required=True, help='Directory where model files will be stored')
    parser.add_argument('--images_dir', required=True, help='Directory where the images are stored')
    parser.add_argument('--labels_dir', required=True, help='Directory where the labels are stored')
    args = parser.parse_args()
    """
    
    python yaml_file.py --model_dir /path/to/model_dir --images_dir /path/to/images_dir --labels_dir /path/to/labels_dir
    
    """
    #create and save the yaml file into the data folder
    create_yaml(args.model_dir, args.images_dir, args.labels_dir)

