"""
This file contains methods for writing and reading csv files applicable for other methods used by library eg.visualisation
Look at main function for usage example


"""


import pandas as pd

def save_csv_w_annot(filepath,annot,annot_index,output_data):
    delimiter="\t"
    annotations = pd.DataFrame(annot,index=annot_index,dtype=object)
    df = pd.DataFrame(output_data)
    annotations.to_csv(filepath,sep=delimiter)
    df.to_csv(filepath,sep=delimiter,mode="a")

def read_csv_w_annot(filepath):
    delimiter="\t"
    found_end=False
    n=1
    while found_end==False:
        annotations = pd.read_csv(filepath,sep=delimiter,nrows=n,dtype=object,index_col=0)
        n=n+1
        if annotations.T.iloc[0,-1] == "########end":
            found_end=True
        elif n==20:
            print("Annotations not found")
            annotations=None
            read_data = pd.read_csv(filepath,sep=delimiter,skiprows=None,index_col=0)
            return annotations,read_data
            
    read_data = pd.read_csv(filepath,sep=delimiter,skiprows=n,index_col=0)
    return annotations,read_data

def get_groups(annot):
    group1 = annot.T.group1.dropna().tolist()
    group2 = annot.T.group2.dropna().tolist()
    return group1,group2
    
def read_csv_any(filepath):
    """
    a1 - annotation df
    a2 - beta values df
    """
    
    a1,a2 = read_csv_w_annot(filepath)
    
    if isinstance(a1,pd.DataFrame):
        gra,grb = get_groups(a1)
        group1_data = a2[gra]
        group2_data = a2[grb]
        return group1_data,group2_data
    return a2
 
def main():
    import os 
    
    #data for annotation part of csv containing groups column names etc
    groupA = ["a","b","c"]
    groupB = ["d","e"]
    
    #end point - for longer annotation part
    theend = ["########end"]
    
    
    _annot_index = ["group1","group2","end"] #must end with end

    #annotation list and data example
    _annot = [groupA,groupB,theend]
    _output_data = {
        'a':[1,2,3], 
        'b':[1,2,3],
        'c':[5,5,5],
        'd':[6,6,7],
        'e':[3,4,5]        
    }

    filepath = os.path.join(os.getcwd(),"csv_with_annotations2.csv")
        
    save_csv_w_annot(filepath,_annot,_annot_index,_output_data)
    output = read_csv_any(filepath)
    return output
    

if __name__ == '__main__':
    
    
    output = main()