# Prepare sample sheet for demultiplexing. 

# The sample sheet generated by Clarity LIMS is not compatible with
# bcl2fastq2. This script adds two headers to the sample sheet to make
# it readable by bcl2fastq2.

# Sample sheet management:
#  * LIMS mode *
#  - This script takes the sample sheet from the "Input sample sheet" ResultFile 
#    object associated with the task's LIMS process. If there is no sample sheet, 
#    it takes it from <run folder>/SampleSheet.csv
#  - It stores the *input* sample sheet in the LIMS process if it was fetched from
#    somewhere else, but this is not used further.
#  - It adds the headers and copies it to the process in the LIMS, on the 
#    "Demultiplexing sample sheet" ResultFile object. 
#  - Other tasks always take the sample sheet from the LIMS process ("Demultiplexing
#    sample sheet"), not any local file. The demultiplexing process writes it to
#    DemultiplexingSampleSheet.<process_id>.csv because bcl2fastq2 needs an on-disk copy.
#
#  * Command line mode *
#  - This script reads the sample sheet from SampleSheet.csv and writes it to
#    DemutliplexingSampleSheet.csv. The input sample sheet file may be overridden
#    with the --input-sample-sheet option.
#  - Other tasks read the sample sheet from DemultiplexingSampleSheet.csv or 
#    another file provided on the command line.

# Since this script must deal with two sample sheets, input and output, it can't
# use the "taskmgr.Task" interface, and has to manually I/O the sample sheets.

import os.path, sys
import argparse
import subprocess
import datetime
import re

from genologics.lims import *
from common import nsc, utilities, taskmgr

TASK_NAME = "20. Prepare SampleSheet"
TASK_DESCRIPTION = "Prepare sample sheet for bcl2fastq2"
TASK_ARGS = ['work_dir', 'lanes']


def main(task):
    """Sample sheet processing."""

    task.add_argument(
            '--input-sample-sheet',
            default=None, # would use work-dir/SampleSheet.csv but dont know work-dir
            help="Path to source sample sheet to transform"
            )
    task.add_argument(
            '--output-sample-sheet',
            default=None,
            help="Path to destination sample sheet to write"
            )
    task.running()

    # Flag to track if the sample sheet was found right on the current LIMS process
    # If the sample sheet was found on the current process, there is no need to 
    # re-upload it.
    found_on_process = False
    sample_sheet = None
    sample_sheet_lims_artifact = None # For use when uploading the input sample sheet
    if task.process:
        # Try to get the input sample sheet:
        
        # On the current LIMS process
        for o in task.process.all_outputs(unique=True):
            if o.output_type == "ResultFile" and o.name == nsc.INPUT_SAMPLE_SHEET:
                sample_sheet_lims_artifact = o
                if len(o.files) == 1:
                    sample_sheet = o.files[0].download()
                    found_on_process = True


        # From a file (fall through to non-LIMS case)
        if not sample_sheet:
            sample_sheet_path = os.path.join(task.work_dir, "SampleSheet.csv")

    else:
        sample_sheet_path = task.args.input_sample_sheet
        if not sample_sheet_path:
            sample_sheet_path = os.path.join(task.work_dir, "SampleSheet.csv")
            if not os.path.exists(sample_sheet_path):
                sample_sheet_path = os.path.join(task.bc_dir, "SampleSheet.csv")
        

    # If not using LIMS, or if LIMS sources failed, try to get a file
    if not sample_sheet:
        try:
            sample_sheet = open(sample_sheet_path, 'rb').read()
        except IOError:
            task.fail("Couldn't find the sample sheet")

    # Upload the input sample sheet as a record if -using LIMS, -it wasn't there already
    if task.process and not found_on_process:
        utilities.upload_file(
                task.process,
                name = nsc.INPUT_SAMPLE_SHEET,
                path = "SampleSheet.csv",
                data = sample_sheet
                )

    # Always replace special characters with ? and fix line breaks
    sample_sheet = replace_special_chars(sample_sheet)

    # Doctor the sample sheet, only if using HiSeq and it doesn't have [Data] header
    instrument = utilities.get_instrument_by_runid(task.run_id)
    if instrument.startswith("hiseq") and not re.search(r"\[Data\]", sample_sheet):
        sample_sheet = convert_from_bcl2fastqv1(sample_sheet)
    # NextSeq run will have index2 reversed by sample sheet generator -- unconditionally.
    # However, single-read flow cells read the index2 in forward direction. So we re-reverse
    # it.
    if instrument == "nextseq":
        if task.process: # LIMS mode
            seq_process = utilities.get_sequencing_process(task.process)
            if seq_process and seq_process.udf.get('Read 1 Cycles') and not seq_process.udf.get('Read 2 Cycles'):
                sample_sheet = reverse_complement_index2(sample_sheet)
        else:
            task.info("Not in LIMS mode, don't know if we should reverse index2, leaving it as on Sample Sheet.")
    # NovaSeq & MiSeq sample sheet generators will replace some invalid characters with _. Including -
    # which is copiously used in our conventions. We bluntly translate all _ back to -, without
    # knowing for sure if they started as - or not.
    if instrument in ["miseq", "novaseq"]:
        sample_sheet = replace_underscores(sample_sheet)

    if task.lanes:
        sample_sheet = filter_lanes(sample_sheet, task.lanes)

    # Post the result, as appropriate...
    if task.process:
        utilities.upload_file(
                task.process,
                name = nsc.SAMPLE_SHEET,
                path = "DemultiplexingSampleSheet.csv",
                data = sample_sheet
                )
    else:
        if task.args.output_sample_sheet:
            path = os.path.join(task.work_dir, task.args.output_sample_sheet)
        else:
            path = os.path.join(task.work_dir, "DemultiplexingSampleSheet.csv")

        open(path, 'wb').write(sample_sheet)
    
    task.success_finish()


def replace_special_chars(sample_sheet_data):
    question_marks = "".join(c if ord(c) < 128 else '?' for c in sample_sheet_data)
    normal_line_breaks = question_marks.replace("\r\r\n", "\n")
    return normal_line_breaks


def rev_comp(sequence):
    COMPLEMENTARY = {'A':'T', 'T':'A', 'G':'C', 'C':'G'}
    return "".join(COMPLEMENTARY[b] for b in reversed(sequence))


def reverse_complement_index2(original_data):
    original_lines = [l.strip("\r\n") for l in original_data.splitlines()]
    data_start = next(i for i, d in enumerate(original_lines) if re.match(r"\[Data\],*$", d))
    data = [l.split(",") for l in original_lines[data_start+1:] if l.strip() != ""]
    try:
        index2_col = next(i for i, c in enumerate(data[0]) if c.lower() == "index2")
        for row in data[1:]:
            #row[index2_col] = str(Seq(row[index2_col]).reverse_complement())
            row[index2_col] = rev_comp(row[index2_col])

        return "\r\n".join(
                original_lines[0:data_start+1] +\
                [",".join(cells) for cells in data]
                )
    except StopIteration: # index2 column not found
        return original_data


def replace_underscores(original_data):
    original_lines = [l.strip("\r\n") for l in original_data.splitlines()]
    data_start = next(i for i, d in enumerate(original_lines) if re.match(r"\[Data\],*$", d))
    data = [l.split(",") for l in original_lines[data_start+1:] if l.strip() != ""]
    try:
        sample_col = next(i for i, c in enumerate(data[0]) if c.lower() == "sample_name")
        project_col = next(i for i, c in enumerate(data[0]) if c.lower() == "sample_project")
        for row in data[1:]:
            row[sample_col] = row[sample_col].replace("_", "-")
            row[project_col] = row[project_col].replace("_", "-")

        return "\r\n".join(
                original_lines[0:data_start+1] +\
                [",".join(cells) for cells in data]
                ) + "\r\n"
    except StopIteration: # sample name or project name column not found
        return original_data


def convert_from_bcl2fastqv1(original_data):
    main_headers = "[Settings]\r\n[Data]\r\n"
    data = [l.strip("\r\n").split(",") for l in original_data.splitlines()]

    replacements = {
            "sampleid": "Sample_Name",
            "description": "Sample_ID",
            "sampleproject": "Sample_Project"
            }


    # Index of Description column:
    description_index = next(i for i, c in enumerate(data[0]) if c.lower() == "description")
    any_empty_description = any(r[description_index] == "" for r in data[1:])
    if any_empty_description:
        del replacements["description"]
        replacements["sampleid"] = "Sample_ID"

    # Replace column names, get index (numeric offset) of the "Index" column
    index_column_index = None
    sample_id_index = None
    for i, cell in enumerate(data[0]):
        rep = replacements.get(cell.lower())
        if rep:
            data[0][i] = rep

        if cell.lower() == "index":
            index_column_index = i
        if cell.lower() == "sampleid":
            sample_id_index = i


    # Add the Sample_Name column if "Description" can't be used as Sample_ID
    # The we use SampleID as sample name and sample ID
    if any_empty_description:
        data[0].append("Sample_Name")
        for r in data[1:]:
            r.append(r[sample_id_index])


    # Split the index into first and second part
    use_dual_index = any(row[index_column_index].find("-") != -1 for row in data[1:])
    if use_dual_index:
        data[0].insert(index_column_index+1, "index2")
        for row in data[1:]:
            indexes = row[index_column_index].split("-")
            if len(indexes) == 1:
                indexes.append("")
                if indexes[0] == "NoIndex":
                    indexes[0] = ""
            row[index_column_index] = indexes[0]
            row.insert(index_column_index+1, indexes[1])
    else:
        for row in data[1:]:
            if row[index_column_index] == "NoIndex":
                row[index_column_index] = ""

    return main_headers + "\r\n".join(
            [",".join(cells) for cells in data]
            )
    

def filter_lanes(original_data, lanes):
    data = [l.strip("\r\n").split(",") for l in original_data.splitlines()]
    filtered = []
    lane_col_index = -1
    data_section = False
    for row in data:
        if data_section:
            for i, col_header in enumerate(row):
                if col_header.lower() == "lane":
                    lane_col_index = i
            data_section = False
            filtered.append(row)
        elif lane_col_index == -1 and row[0].lower() == "[data]":
            data_section = True
            filtered.append(row)
        elif lane_col_index == -1:
            filtered.append(row)
        else:
            try:
                if int(row[lane_col_index]) in lanes:
                    filtered.append(row)
            except (IndexError, ValueError):
                filtered.append(row)
    return "\r\n".join([",".join(cells) for cells in filtered]) + "\r\n"


if __name__ == "__main__":
    with taskmgr.Task(TASK_NAME, TASK_DESCRIPTION, TASK_ARGS) as task:
        main(task)
