#!python

""" This script takes the JSON file output by hivnetworkcsv and adds the
data dictionary for subject attributes to create an annotated JSON suitable
for rendering by hivtrace-viz (https://github.com/veg/hivtrace-viz)

Author:
    Sergei L Kosakovsky Pond (spond@temple.edu)

Version:
    v1.0.0 (2017-05-09)
    v1.1.0 (2017-09-04) : added support for 'enums', fixed bugs

Example:
    python3 scripts/hivnetworkannotate -t examples/lanl.tsv -n examples/network.json
        -f "Country" "Country" "String" "x: 'N/A' f x=='-' else x"
        -f "Year" "Year" "Number" "x: 'N/A' if x=='-' else x"
        -f "RiskFactor" "Risk factor" "enum" '{"SG" : "homosexual", "SB" : "bisexual", "SM" : "male sex with male", "SH" : "heterosexual", "SW" : "sex worker", "SU" : "sexual transmission, unspecified type", "PH" : "hemophiliac", "PB" : "Blood transfusion", "PI" : "IV drug use", "MB" : "Mother-baby", "NO" : "Nosocomial", "EX" : "Experimental", "NR" : "not recorded (or unknown)", "OT" : "other", "-" : "N/A"}'

    See https://github.com/veg/hivclustering/wiki/hivnetworkannotate

"""

import csv
import argparse
import sys
import json
import re
import datetime
import functools
from hivclustering.networkbuild import *

def enum_mapper (key, dict):
    try:
        return dict[key]
    except KeyError:
        return None

def ensure_key (d,key,value = None):
    if key not in d:
        if value is None:
            d[key] = {}
        else:
            d[key] = value


    return d[key]

# Construct keys using specified fields and delimiter from keying field in schema 
def construct_key(record, fields, delimiter):
    try:
        return delimiter.join([str(record[field]) for field in fields])
    except KeyError as e:
        print(f"Missing field {e} in record {record}")
        return None

def construct_node_key_from_schema(node_id, schema_key_fields, delimiter):
    node = dict(zip(schema_key_fields, node_id.split(delimiter)))
    key_parts = []
    for field in schema_key_fields:

        # Assume each node's 'id' contains a delimiter-separated string of fields
        # Extract the value corresponding to each field
        field_value = node.get(field, None)

        if field_value is not None:
            key_parts.append(str(field_value))
        else:
            print(f"Field '{field}' not found in node: {node}")
            return None  # Return None if any key field is missing
    return delimiter.join(key_parts)

#-------------------------------------------------------------------------------


arguments = argparse.ArgumentParser(description='Annoated the JSON file representing the PIRC network with data attributes from clinical data and isolation dates.')

arguments.add_argument  ('-o', '--output', help    = 'Output the annotated JSON network file to', nargs = '?', type = argparse.FileType('w'), default = sys.stdout)
arguments.add_argument  ('-n', '--network',   help    = 'The input network file to process', nargs = '?', type = argparse.FileType('r'), default = sys.stdin)
arguments.add_argument  ('-x', '--missing', help  = 'If desired, provide a value to inject for nodes that do not have an attribute value specified', required = False, nargs = 2, action = 'append')
arguments.add_argument  ('-X', '--clear', help = 'Flush existing attributes', required = False, action = 'store_true')
arguments.add_argument  ('-i', '--index',   help    = 'The name of the column that indexes records (patient ID); default is to index on the first column', type = str)
arguments.add_argument  ('-r', '--inplace', help = 'Write attributes to the input file (cannot be stdin)', required = False, action = 'store_true')
arguments.add_argument  ('-F', '--full-inject', help = 'Add nodes from the attributes file that are not already present in the network', required = False, action = 'store_true')

input_group = arguments.add_mutually_exclusive_group(required=True)
input_group.add_argument  ('-a', '--attributes',   help    = 'The JSON file with node attributes', type = argparse.FileType('r'))
input_group.add_argument  ('-t', '--tab',   help    = 'A TSV file with node attributes', type = argparse.FileType('r'))
input_group.add_argument  ('-c', '--csv',   help    = 'A CSV file with node attributes', type = argparse.FileType('r'))

field_input_group = arguments.add_mutually_exclusive_group(required=True)
field_input_group.add_argument ('-f', '--field', help = 'Describe an argument to be added to invididual nodes as "name" "label" "type" "transform"; currently supported types are "String", "enum", "Date", "Number"; transform must be specified as a lambda, an empty string to use an identity map, or a python style dict to specify an enum; "fulldate" is a predefined option to reformat the date using the default hivtrace-viz format', nargs = 4, action = 'append')
field_input_group.add_argument ('-g', '--fields-file', help = 'Read in fields from a JSON file', type = argparse.FileType('r'))

import_settings = arguments.parse_args()

results_json = json.load (import_settings.network)

root_trace_results = False

if('trace_results' in results_json.keys()):
    root_trace_results = True

network_json    = ht_process_network_json(results_json)

# set up record filtering

network_attribute_key = "patient_attribute_schema"
node_attribute_key    = "patient_attributes"
inject_missing_value  = import_settings.missing

key_fields = ['ehars_uid'];
key_delimiter = '~';

ensure_key (network_json, network_attribute_key)

field_transformations = {}
field_names = {}
predefined_transforms = {'YYYYMMDD' : 's: datetime.datetime.strptime (s,"%Y-%m-%d").strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]+"Z"',
                         'YYYY' : 's: datetime.datetime.strptime (s,"%Y").strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]+"Z"'}

uninjected_set      = {}

if import_settings.field:
    field_settings = import_settings.field
else:
    # If fields file supplied, get fields from that first
    fields_file_json = json.load(import_settings.fields_file)
    for key, d in fields_file_json.items():
        ensure_key(d, 'enum')

    # Check for _keying_ value
    key_fields = fields_file_json.get('keying', {}).get('fields', ['ehars_uid'])
    key_delimiter = fields_file_json.get('keying', {}).get('delimiter', '~')

    #Filter out items that don't have a label
    fields_file_json = {k : v for k,v in fields_file_json.items() if "label" in v.keys()}
    field_settings = [[k, v["label"], v["type"], v.get("transform", "")] for k,v in fields_file_json.items()]


for key_pair in field_settings:
    network_json [network_attribute_key][key_pair[1]] = {'name' : key_pair[1], 'type': key_pair[2], 'label' : key_pair[1]}
    field_names [key_pair[0]] = key_pair[1]
    uninjected_set [key_pair[1]] = set ()
    if key_pair[2] == 'enum':
        mapping = eval (key_pair[3])
        network_json [network_attribute_key][key_pair[1]]['type'] = "String"
        network_json [network_attribute_key][key_pair[1]]['enum'] = list (mapping.values())
        field_transformations [key_pair[0]] = functools.partial (enum_mapper, dict = mapping)
    else:
        if len (key_pair[3]) == 0 :
            field_transformations [key_pair[0]] = lambda x: x
        elif key_pair[3] in predefined_transforms:
            field_transformations [key_pair[0]] = eval ("lambda " + predefined_transforms[key_pair[3]])
        else:
            field_transformations [key_pair[0]] = eval ("lambda " + key_pair[3])

total_records = 0

if import_settings.attributes: # JSON input
    to_import_raw = json.load (import_settings.attributes)
    to_import = {}
    for v in to_import_raw:
        key = construct_key(v, key_fields, key_delimiter)
        if key:
            if key not in to_import:
                to_import[key] = []
            to_import[key].append(v)
    total_records = len(to_import_raw)
else: # TSV/CSV import
    if import_settings.tab:
        csv_reader = csv.reader (import_settings.tab, delimiter = '\t')
    else:
        csv_reader = csv.reader (import_settings.csv, delimiter = ',')

    fields = next (csv_reader)

    index_on = 0
    if import_settings.index:
        idx = fields.index (import_settings.index)
        if idx < 0:
            raise "Invalid field to index on (%s)" % import_settings.index
        index_on = lambda L : L[idx]
    else:
        if import_settings.fields_file:
             indices = []
             for f in key_fields:
                indices.append (fields.index (f))
                if indices[-1] < 0:
                    raise "Invalid field to index on (%s)" % f
             index_on = lambda L : key_delimiter.join ([L[k] for k in indices])
    
    to_import = {}

    try:
        for line in csv_reader:
            line_index = index_on (line)
            record = {}
            for i, k in enumerate (line):
                record[fields[i]] = k
            if line_index not in to_import:
                to_import [line_index] = []
            to_import [line_index].append(record)
            total_records += 1
    except IndexError as e:
        print (line, fields)
        raise

id_mapper = lambda x: x

nodes_indexed_by_id = {}
nodes_by_full_id = {}
# Index nodes by their constructed key, rather than their id
for n in network_json["Nodes"]:
    node_key = construct_node_key_from_schema(n['id'],  key_fields, key_delimiter)
    if node_key in nodes_indexed_by_id:
        nodes_indexed_by_id[node_key].append(n)
    else:
        nodes_indexed_by_id[node_key] = [n]
        
    nodes_by_full_id[n['id']] = n
    for f, s in uninjected_set.items():
        s.add(n['id'])  # Track uninjected nodes
    if import_settings.clear:
        if node_attribute_key in n:
            del n[node_attribute_key]  # Clear attributes if requested

def compute_match_score(record, sequence_tokens, key_fields):
    score = 0
    for k, val in record.items():
        if k not in key_fields:
            val_str = str(val) if val is not None else ""
            if val_str:
                val_tokens = re.split(r'[|~]', val_str)
                for token in val_tokens:
                    if token and token in sequence_tokens:
                        score += 1
    return score

def check_has_sequence(record):
    seq_str = record.get('predq_clean_seq', '').upper()
    if seq_str and seq_str not in ('', 'MISSING'):
        return True
    seq_len = record.get('predq_seq_len_clean', '0')
    try:
        if int(seq_len) > 0:
            return True
    except ValueError:
        pass
    return False

injected_nodes = []

# Now match nodes by the key constructed from the input data
for node_key, records in to_import.items():
    if node_key in nodes_indexed_by_id:  # Match by the constructed key
        for n_object in nodes_indexed_by_id[node_key]:
            node_dict = ensure_key(n_object, node_attribute_key)
            
            # Select the best record from the list of records for this node
            node_id_tokens = set(re.split(r'[|~]', n_object['id']))
            key_tokens = set(re.split(r'[|~]', node_key))
            sequence_tokens = node_id_tokens - key_tokens
            
            best_record = None
            max_score = -1
            # Check for sequence-specific identifiers matching the node ID tokens
            for record in records:
                score = compute_match_score(record, sequence_tokens, key_fields)
                if score > max_score:
                    max_score = score
                    best_record = record
            
            if best_record is None:
                best_record = records[0]
                
            for k, val in best_record.items():
                if k in field_transformations:
                    store_this = field_transformations[k](val)
                    if store_this is not None and node_dict is not None:
                        node_dict[field_names[k]] = store_this
                        if n_object['id'] in uninjected_set[field_names[k]]:
                            uninjected_set[field_names[k]].remove(n_object['id'])
            
            if node_dict is not None:
                node_dict['has_sequence'] = True
                node_dict['sequence_status'] = 'TNA-eligible'
                if 'poor_quality' in best_record:
                    try:
                        node_dict['poor_quality'] = int(best_record['poor_quality'])
                    except ValueError:
                        pass
                if 'fraction_ambig' in best_record:
                    try:
                        node_dict['fraction_ambig'] = float(best_record['fraction_ambig'])
                    except ValueError:
                        pass
                if 'TNA_eligible' in best_record:
                    try:
                        node_dict['TNA_eligible'] = int(best_record['TNA_eligible'])
                    except ValueError:
                        pass
    elif import_settings.full_inject:  # If full-inject is enabled, construct a new node
        # For injected nodes, choose the first record (all records are per-sequence)
        best_record = records[0]
            
        new_node = {
            "id": node_key,
            node_attribute_key: {}
        }
        node_dict = new_node[node_attribute_key]
        for k, val in best_record.items():
            if k in field_transformations:
                store_this = field_transformations[k](val)
                if store_this is not None:
                    node_dict[field_names[k]] = store_this
        
        has_seq = check_has_sequence(best_record)
        node_dict['has_sequence'] = has_seq
        node_dict['sequence_status'] = 'TNA-eligible' if has_seq else 'none'
        
        if has_seq:
            if 'poor_quality' in best_record:
                try:
                    node_dict['poor_quality'] = int(best_record['poor_quality'])
                except ValueError:
                    pass
            if 'fraction_ambig' in best_record:
                try:
                    node_dict['fraction_ambig'] = float(best_record['fraction_ambig'])
                except ValueError:
                    pass
            if 'TNA_eligible' in best_record:
                try:
                    node_dict['TNA_eligible'] = int(best_record['TNA_eligible'])
                except ValueError:
                    pass
        
        injected_nodes.append(new_node)

if import_settings.full_inject and len(injected_nodes) > 0:
    # 1. Append newly injected nodes
    network_json["Nodes"].extend(injected_nodes)
    
    # 2. Tag each node with its pre-sort index
    for i, node in enumerate(network_json["Nodes"]):
        node["__temp_index__"] = i
    
    # 3. Sort all nodes alphabetically by ID to take advantage of Front (prefix) compression
    network_json["Nodes"] = sorted(network_json["Nodes"], key=lambda x: x["id"])
    
    # 4. Build mapping of old index -> new sorted index
    old_to_new_index = {node["__temp_index__"]: i for i, node in enumerate(network_json["Nodes"])}
    
    # 5. Clean up temporary index tag
    for node in network_json["Nodes"]:
        del node["__temp_index__"]
    
    # 6. Remap Edges source/target indices to match the new sorted node order
    for edge in network_json.get("Edges", []):
        edge["source"] = old_to_new_index[edge["source"]]
        edge["target"] = old_to_new_index[edge["target"]]

print ("\nImport summary", file = sys.stderr)
print ("\t Records in file  : %d" % total_records, file = sys.stderr)
print ("\t Nodes in network : %d" % len(network_json ["Nodes"]), file = sys.stderr)

for k, v in uninjected_set.items():
    print ("\t Field '%s'  : %d records imported" % (k, len(network_json ["Nodes"]) -len (v)), file = sys.stderr)
    

print ("\n", file = sys.stderr)
    
if inject_missing_value:
    for values in inject_missing_value:        
        for node_id in uninjected_set[values[0]]:
            node_dict = ensure_key (nodes_by_full_id[node_id], node_attribute_key)
            node_dict[values[0]] = values[1]


if import_settings.inplace and import_settings.network is not sys.stdin:
    import_settings.output = open (import_settings.network.name, "w")
    
if network_attribute_key in network_json:
    if 'has_sequence' not in network_json[network_attribute_key]:
        network_json[network_attribute_key]['has_sequence'] = {
            'name': 'has_sequence',
            'type': 'Boolean',
            'label': 'Has Sequence'
        }
    if 'sequence_status' not in network_json[network_attribute_key]:
        network_json[network_attribute_key]['sequence_status'] = {
            'name': 'sequence_status',
            'type': 'String',
            'label': 'Sequence Status'
        }
    if 'poor_quality' not in network_json[network_attribute_key]:
        network_json[network_attribute_key]['poor_quality'] = {
            'name': 'poor_quality',
            'type': 'Number',
            'label': 'poor_quality'
        }
    if 'fraction_ambig' not in network_json[network_attribute_key]:
        network_json[network_attribute_key]['fraction_ambig'] = {
            'name': 'fraction_ambig',
            'type': 'Number',
            'label': 'fraction_ambig'
        }
    if 'TNA_eligible' not in network_json[network_attribute_key]:
        network_json[network_attribute_key]['TNA_eligible'] = {
            'name': 'TNA_eligible',
            'type': 'Number',
            'label': 'TNA_eligible'
        }

for n in network_json["Nodes"]:
    node_dict = ensure_key(n, node_attribute_key)
    if 'has_sequence' not in node_dict:
        node_dict['has_sequence'] = True
    if 'sequence_status' not in node_dict:
        node_dict['sequence_status'] = 'TNA-eligible'

if 'Settings' in network_json and 'compact_json' in network_json['Settings'] and network_json['Settings']['compact_json']:
     ht_compress_network_json (network_json)

# Return with trace_results key
to_return = {}
if(root_trace_results):
    to_return = results_json.copy()  # Preserve original top-level keys
    to_return["trace_results"] = network_json  # Update trace_results with processed data
else:
    to_return = network_json


json.dump(to_return, import_settings.output, indent=2, separators=(',', ':'))
