# This code was developed and authored by Jerzy Twarowski in Malkova Lab at the University of Iowa 
# Contact: jerzymateusz-twarowski@uiowa.edu, tvarovski1@gmail.com

import regex as re

def createMMBSearchReference(file_path_in: str, file_path_out: str) -> None:
    """
    Reads a FASTA file containing genome data and creates a new file with a modified header line for each chromosome.
    It's purpose is to create a FASTA file that can be used as a reference for MMBSearch.

    Args:
        file_path_in (str): The path to the input FASTA file.
        file_path_out (str): The path to the output file to be created.

    Returns:
        None
    """
    #good for human genome
    file_in = open(file_path_in, "r")
    save=False
    readlines=True
    while readlines:
        try:
            line = file_in.readline()
        except:
            print("Failed to read line. EOF? Exiting...")
            readlines=False
            break

        if line[0] == ">":
            if re.search("chromosome.*Primary.Assembly$", line) != None:
                print(line)
                linewords = line.split()
                chromosome=linewords[4].strip(",")
                if (chromosome=="1") | (chromosome=="2"):
                    chromosome=f">chr0{chromosome}\n"
                elif chromosome=="X":
                    chromosome=f">chrX\n"
                elif chromosome=="Y":
                    chromosome=f">chrY\n"
                else:
                    chromosome=f">chr{chromosome}\n"
                line=chromosome
                print(line)
                save=True
            #fix for mitochondrial chr
            elif re.search("mitochondrion, complete genome$", line) != None:
                print(line)
                line=f">chrM\n"
                print(line)
                save=True
            else:
                save=False

        if save:
            with open(file_path_out, "a") as file_out:
                file_out.write(line)