##############################################################################
# Import necessary modules

import os
import sys
import argparse
import csv
from collections import defaultdict
import random


##############################################################################

current_version = '0.8'


##############################################################################
# Parses input using argparse module

# Initiates argument parser
parser = argparse.ArgumentParser(description='Make order file')

# Non-optional arguments: List of genes, screen type and name of output file

parser.add_argument('order_folder', type=str)

parser.add_argument('-t', '--test', action='store_true')
parser.add_argument('-x', '--extra', type=str, nargs=2)
parser.add_argument('-x2', '--extra2', type=str, nargs=2)
parser.add_argument('-i', '--inter', type=int, default=5)
parser.add_argument('-l', '--low', type=int, default=500)

# Saves input to args object
args = parser.parse_args()


##############################################################################
# 

directiory, chip_name = os.path.split(args.order_folder[:-1])
legend_file = os.path.join(args.order_folder, chip_name + '_legend.csv')
order_file = os.path.join(args.order_folder, chip_name + '_order.csv')
align_folder = os.path.join(args.order_folder, 'Align')


##############################################################################
# Locate all align files

lib_names = []

with open(legend_file, 'r') as legend_open:
    legend_csv = csv.reader(legend_open)

    for line in legend_csv:
        lib_name = line[0]
        lib_names.append(lib_name.split('_align')[0])

align_list = []

for root, dirs, files in os.walk(align_folder):
    for fil in files:

        if '_align.csv' in fil:

            full_path = os.path.join(root, fil)
            name = os.path.split(full_path)[-1][:-10]

            found = False

            for lib_name in lib_names:
                if lib_name == name:

                    if not found:
                        found = True

                    else:
                        sys.exit('Error: Multiple legend entries found\n' + lib_name + '\n' + full_path)

            if found:
                align_list.append(full_path)
            else:
                sys.exit('Error: No legend entry found:\n' + name)


##############################################################################
# For each align file, remove duplicate and write order file

outputs = []

for align_file in align_list:

    data = []
    copies = 1

    with open(align_file, 'r') as align_open:  
           
        align_csv = csv.reader(align_open, delimiter=',')

        for line in align_csv:
            data.append(line)

    #
    data_filter = dict((x[1], x) for x in data).values()
    data_filter.sort(key=lambda x: x[0])

    name = align_file.split('/')[-1].split('.')[-2][:-6]
    #order_file = os.path.join(args.order_folder, name) + '_order.csv'

    if 'hU6' in name or 'mU6' in name or 'inter' in name:
        copies = args.inter

    if args.extra:
        extr_string, extr_copy = args.extra

        if extr_string in name:
            copies = int(extr_copy)

    if args.extra2:
        extr_string, extr_copy = args.extra2

        if extr_string in name:
            copies = int(extr_copy)

    removed = len(data) - len(data_filter)
    #print(order_file)

    while copies * len(data_filter) < args.low:
        copies += 1

    print(str(removed) + ' duplicate elements removed from ' + name)
    print(str(copies) + ' copies of ' + name + ' added.')

    outputs.append((copies, data_filter))

##############################################################################
#

total = 0

for copies, output in outputs:

    total += copies * len(output)

print('Total elements: ' + str(total))

if args.test:
    sys.exit('Warning: Files not written')

with open(order_file, 'w') as order_open:
    order_csv = csv.writer(order_open, delimiter=',', lineterminator='\n')

    for copies, output in outputs:

        for t in range(copies):

            for line in output:
                order_csv.writerow(line)


##############################################################################