#########################################################################################
# Author: <NAME>
# Date Started: 2016-07-26
# Environment: Python3
# License: See LICENSE
# Purpose: Load dataset and create interfaces for piping the data to the model.
##########################################################################################

import csv
import numpy as np

import lib_paths
import atchley_factors as vector_representation


def _load_repertoires(data_dir):
  repertoires = dict()
  with open(data_dir+'/data/answers.csv', 'r') as keyfile_stream:
    keyfile_reader = csv.DictReader(keyfile_stream, delimiter=',')
    for keyfile_row in keyfile_reader:
      sample_id = keyfile_row['Sample']
      diagnosis_id = keyfile_row['Diagnosis']
      sequences = dict()
      path = data_dir+'/data/sample_'+keyfile_row['Sample']+'.csv'
      with open(path, 'r') as sample_stream:
        sample_reader = csv.DictReader(sample_stream, delimiter=',')
        for sample_row in sample_reader:
          sequence = sample_row['Sequence']
          count = float(sample_row['Count'])
          if sequence not in sequences:
            sequences[sequence] = count
          else:
            sequences[sequence] += count
      repertoires[sample_id] = {
        'Diagnosis': diagnosis_id,
        'Sequences': sequences
      }
  return repertoires


def load_repertoires(data_file):
    repertoires = dict()
    seqs_per_ind = dict()
    with open(data_file, 'r') as fh:
        for l in fh:
            seq, sample_id, diagnosis_id = l.strip().split()
            if sample_id not in repertoires:
                repertoires[sample_id] = {
               'Diagnosis': diagnosis_id,
               'Sequences': dict()
            }
            count = 1.0  # So the true count is actually not used. The count is only used as an indication of zero or more.
            if seq not in repertoires[sample_id]['Sequences']:
                repertoires[sample_id]['Sequences'][seq] = count
            else:
                repertoires[sample_id]['Sequences'][seq] += count
    return repertoires


def process_repertoires(repertoires, snip_size=6):
  repertoires_snip = {}
  for sample, repertoire in repertoires.items():
    diagnosis = repertoire['Diagnosis']
    snips = {}
    for sequence, count in repertoire['Sequences'].items():
      stop = len(sequence)-snip_size+1
      for i in range(stop):
        snip = sequence[i:i+snip_size]
        if snip not in snips:
          snips[snip] = count
        else:
          snips[snip] += count
    repertoires_snip[sample] = {
      'Diagnosis': diagnosis,
      'Snips': snips
    }

  num_samples = len(repertoires)
  max_snips = -1
  num_features = snip_size*vector_representation.length

  for sample, repertoire in repertoires_snip.items():
    num_snips = len(repertoire['Snips'])
    if num_snips > max_snips:
      max_snips = num_snips

  xs = np.zeros((num_samples, max_snips, num_features), dtype=np.float32)  # Features
  cs = np.zeros((num_samples, max_snips), dtype=np.float32)        # Snippet count
  ys = np.zeros((num_samples), dtype=np.float32)  # Labels

  for i, (sample, repertoire) in enumerate(sorted(repertoires_snip.items(), key=lambda item: item[0])):
    for j, (snip, count) in enumerate(repertoire['Snips'].items()):
      xs[i,j,:] = vector_representation.features(snip)
      cs[i,j] = float(count)
    ys[i] = float(repertoire['Diagnosis'])
  return xs, cs, ys


#########################################################################################
# Author: <NAME>
# Date Started: 2016-07-26
# Environment: Python3
# License: See LICENSE
# Purpose: Load dataset and create interfaces for piping the data to the model.
##########################################################################################

import csv
import numpy as np

import lib_paths
import atchley_factors as vector_representation


def _load_repertoires(data_dir):
  repertoires = dict()
  with open(data_dir+'/data/answers.csv', 'r') as keyfile_stream:
    keyfile_reader = csv.DictReader(keyfile_stream, delimiter=',')
    for keyfile_row in keyfile_reader:
      sample_id = keyfile_row['Sample']
      diagnosis_id = keyfile_row['Diagnosis']
      sequences = dict()
      path = data_dir+'/data/sample_'+keyfile_row['Sample']+'.csv'
      with open(path, 'r') as sample_stream:
        sample_reader = csv.DictReader(sample_stream, delimiter=',')
        for sample_row in sample_reader:
          sequence = sample_row['Sequence']
          count = float(sample_row['Count'])
          if sequence not in sequences:
            sequences[sequence] = count
          else:
            sequences[sequence] += count
      repertoires[sample_id] = {
        'Diagnosis': diagnosis_id,
        'Sequences': sequences
      }
  return repertoires


def load_repertoires(data_file):
    repertoires = dict()
    seqs_per_ind = dict()
    with open(data_file, 'r') as fh:
        for l in fh:
            seq, sample_id, diagnosis_id = l.strip().split()
            if sample_id not in repertoires:
                repertoires[sample_id] = {
               'Diagnosis': diagnosis_id,
               'Sequences': dict()
            }
            count = 1.0  # So the true count is actually not used. The count is only used as an indication of zero or more.
            if seq not in repertoires[sample_id]['Sequences']:
                repertoires[sample_id]['Sequences'][seq] = count
            else:
                repertoires[sample_id]['Sequences'][seq] += count
    return repertoires


def process_repertoires(repertoires, snip_size=6):
  repertoires_snip = {}
  for sample, repertoire in repertoires.items():
    diagnosis = repertoire['Diagnosis']
    snips = {}
    for sequence, count in repertoire['Sequences'].items():
      stop = len(sequence)-snip_size+1
      for i in range(stop):
        snip = sequence[i:i+snip_size]
        if snip not in snips:
          snips[snip] = count
        else:
          snips[snip] += count
    repertoires_snip[sample] = {
      'Diagnosis': diagnosis,
      'Snips': snips
    }

  num_samples = len(repertoires)
  max_snips = -1
  num_features = snip_size*vector_representation.length

  for sample, repertoire in repertoires_snip.items():
    num_snips = len(repertoire['Snips'])
    if num_snips > max_snips:
      max_snips = num_snips

  xs = np.zeros((num_samples, max_snips, num_features), dtype=np.float32)  # Features
  cs = np.zeros((num_samples, max_snips), dtype=np.float32)        # Snippet count
  ys = np.zeros((num_samples), dtype=np.float32)  # Labels

  for i, (sample, repertoire) in enumerate(sorted(repertoires_snip.items(), key=lambda item: item[0])):
    for j, (snip, count) in enumerate(repertoire['Snips'].items()):
      xs[i,j,:] = vector_representation.features(snip)
      cs[i,j] = float(count)
    ys[i] = float(repertoire['Diagnosis'])
  return xs, cs, ys


######################################################################################### # Author: <NAME> # Date Started: 2016-07-26 # Environment: Python3 # License: See LICENSE # Purpose: Load dataset and create interfaces for piping the data to the model. ########################################################################################## # So the true count is actually not used. The count is only used as an indication of zero or more. # Features # Snippet count # Labels