/*
 *  Wellcome Trust Sanger Institute
 *  Copyright (C) 2011  Wellcome Trust Sanger Institute
 *  
 *  This program is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU General Public License
 *  as published by the Free Software Foundation; either version 2
 *  of the License, or (at your option) any later version.
 *  
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *  
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "parse_phylip.h"
#include "alignment_file.h"

#include <zlib.h>
#include <sys/types.h>
#include "kseq.h"
#include "string_cat.h"

KSEQ_INIT(gzFile, gzread)

int num_samples;
int num_snps;
char ** sequences;
char ** phylip_sample_names;
int * internal_node;
sample_statistics ** statistics_for_samples;


int update_sequence_base(char new_sequence_base, int sequence_index, int base_index)
{
	if(sequences[sequence_index][base_index] != new_sequence_base)
	{
	   sequences[sequence_index][base_index] = new_sequence_base;	
	   return 1;
    }
	return 0;
}

void set_internal_node(int internal_node_value,int sequence_index)
{
	internal_node[sequence_index] = internal_node_value;
}

int get_internal_node(int sequence_index)
{
	return internal_node[sequence_index];
}

void get_sequence_for_sample_name(char * sequence_bases, char * sample_name)
{
    int sequence_index;
    sequence_index = find_sequence_index_from_sample_name(sample_name);
    if(sequence_index < 0)
    {
        printf("Could not find sequence name %s with index %d\n", sample_name,sequence_index);
        exit(1);
    }
    get_sequence_for_sample_index(sequence_bases, sequence_index);
    
}

void get_sequence_for_sample_index(char * sequence_bases, int sequence_index)
{
    memcpy(sequence_bases, sequences[sequence_index], size_of_string(sequences[sequence_index]) +1);
}

void fill_in_unambiguous_gaps_in_parent_from_children(int parent_sequence_index, int * child_sequence_indices, int num_children)
{
	int snp_counter = 0; 
	for(snp_counter = 0; snp_counter < num_snps ; snp_counter++)
	{
		int real_base_found = 0;
		int child_counter = 0;
		for(child_counter = 0 ; child_counter < num_children ; child_counter++)
		{
			int child_index = child_sequence_indices[child_counter];
			if(!( toupper(sequences[child_index][snp_counter]) == 'N' || sequences[child_index][snp_counter] == '-' ))
			{
				real_base_found = 1;
				break;
			}
		}
		
		if(real_base_found == 0 &&  toupper(sequences[parent_sequence_index][snp_counter]) != 'N' && sequences[parent_sequence_index][snp_counter] != '-')
		{
			sequences[parent_sequence_index][snp_counter] = 'N';	
		}
	}
}

void fill_in_unambiguous_bases_in_parent_from_children_where_parent_has_a_gap(int parent_sequence_index, int * child_sequence_indices, int num_children)
{
    int snp_counter = 0;

    for(snp_counter = 0; snp_counter < num_snps ; snp_counter++)
    {
        if(toupper(sequences[parent_sequence_index][snp_counter]) != 'N' && sequences[parent_sequence_index][snp_counter] != '-')
        {
            break;
        }

        int child_counter = 0;
        char comparison_base = '\0';
        for(child_counter = 0 ; child_counter < num_children ; child_counter++)
        {
            int child_index = child_sequence_indices[child_counter];
            if(child_counter == 0)
            {
                comparison_base = toupper(sequences[child_index][snp_counter]);
            }

            if(comparison_base !=  toupper(sequences[child_index][snp_counter])  )
            {
                break;
            }
        }

        if(toupper(sequences[parent_sequence_index][snp_counter]) != comparison_base)
        {
            sequences[parent_sequence_index][snp_counter] = comparison_base;
        }
    }
}

int does_column_contain_snps(int snp_column, char reference_base)
{
	int i;
	reference_base = convert_reference_to_real_base_in_column( snp_column,  reference_base);
	for(i = 0; i < num_samples; i++)
	{
		if(internal_node[i]==1)
		{
			continue;	
		}
		
		if(sequences[i][snp_column] == '\0' || sequences[i][snp_column] == '\n')
		{
			return 0;	
		}
		
		if(sequences[i][snp_column]  != '-' && toupper(sequences[i][snp_column])  != 'N' && sequences[i][snp_column] != reference_base)
		{
			return 1;
		}
	}   
	return 0;
}


char convert_reference_to_real_base_in_column(int snp_column, char reference_base)
{
	int i;
	if(!(reference_base == '-' || toupper(reference_base) == 'N'))
	{
		return reference_base;
	}
	
	for(i = 0; i < num_samples; i++)
	{
		if(sequences[i][snp_column] == '\0' || sequences[i][snp_column] == '\n')
		{
			return reference_base;	
		}
		
		if(sequences[i][snp_column]  != '-' && toupper(sequences[i][snp_column])  != 'N')
		{
			return sequences[i][snp_column];
		}
	}
	return reference_base;
}

int number_of_samples_from_parse_phylip()
{
	return num_samples;
}

sample_statistics ** get_sample_statistics()
{
	return statistics_for_samples;
}


void get_sample_names_from_parse_phylip(char ** sample_names)
{
	int i;
	for(i = 0; i< num_samples; i++)
	{
		sample_names[i] =  (char *) calloc(MAX_SAMPLE_NAME_SIZE,sizeof(char));
		memcpy(sample_names[i], phylip_sample_names[i], size_of_string(phylip_sample_names[i]) +1);
	}
}
	

void filter_sequence_bases_and_rotate(char * reference_bases, char ** filtered_bases_for_snps, int number_of_filtered_snps)
{
	int i,j,reference_index;
	
	for(j = 0; j < number_of_filtered_snps; j++)
	{
		filtered_bases_for_snps[j] = (char *) calloc((num_samples+1),sizeof(char));
	}
		
	for(i = 0; i < num_samples; i++)
	{
		int filtered_base_counter = 0;
		
		for(reference_index = 0; reference_index < num_snps; reference_index++)
		{

			if(reference_bases[reference_index] == '\0')
			{
				break;	
			}
			
			if(reference_bases[reference_index] != '*' && sequences[i][reference_index] != '\0' && sequences[i][reference_index] != '\n')
			{
				filtered_bases_for_snps[filtered_base_counter][i] = sequences[i][reference_index];
				filtered_base_counter++;
			}
		}
	}
	for(j = 0; j < number_of_filtered_snps; j++)
	{
		 filtered_bases_for_snps[j][num_samples] = '\0';	
	}

}


void set_number_of_recombinations_for_sample(int sample_index, int number_of_recombinations)
{
    if( sample_index == -1)
    {
        return;
    }
    ((sample_statistics *) statistics_for_samples[sample_index])->number_of_recombinations = number_of_recombinations;
}


void set_number_of_snps_for_sample(int sample_index, int number_of_snps)
{
    if( sample_index == -1)
    {
        return;
    }
    ((sample_statistics *) statistics_for_samples[sample_index])->number_of_snps = number_of_snps;
}

void set_number_of_blocks_for_sample(char * sample_name,int num_blocks)
{
    int sample_index ;
    sample_index = find_sequence_index_from_sample_name( sample_name);
    if( sample_index == -1)
    {
        return;
    }

    ((sample_statistics *) statistics_for_samples[sample_index])->number_of_blocks = num_blocks;
}



void set_genome_length_without_gaps_for_sample(int sample_index, int genome_length_without_gaps)
{
    if( sample_index == -1)
    {
        return;
    }

    ((sample_statistics *) statistics_for_samples[sample_index])->genome_length_without_gaps = genome_length_without_gaps;
}

void set_genome_length_excluding_blocks_and_gaps_for_sample(int sample_index, int genome_length_excluding_blocks_and_gaps)
{

  if( sample_index == -1)
    {
		return;
	}
	
	((sample_statistics *) statistics_for_samples[sample_index])->genome_length_excluding_blocks_and_gaps = genome_length_excluding_blocks_and_gaps;
}

void set_number_of_branch_bases_in_recombinations(int sample_index, int bases_in_recombinations, int include_gaps)
{

    if( sample_index == -1)
    {
        return;
    }
    
    if (include_gaps > 0) {
      ((sample_statistics *) statistics_for_samples[sample_index])->branch_bases_in_recombinations_including_gaps = bases_in_recombinations;
    } else {
      ((sample_statistics *) statistics_for_samples[sample_index])->branch_bases_in_recombinations = bases_in_recombinations;
    }
}

void set_number_of_bases_in_recombinations(int sample_index, int bases_in_recombinations, int include_gaps)
{
  if( sample_index == -1)
  {
      return;
  }
  
  if (include_gaps > 0) {
    ((sample_statistics *) statistics_for_samples[sample_index])->bases_in_recombinations_including_gaps = bases_in_recombinations;
  } else {
    ((sample_statistics *) statistics_for_samples[sample_index])->bases_in_recombinations = bases_in_recombinations;
  }
  
}




int find_sequence_index_from_sample_name( char * sample_name)
{
	int i;
	
	for(i =0; i< num_samples; i++)
	{
		if(strcmp(sample_name,phylip_sample_names[i]) == 0)	
		{
			return i;
		}
	}
	return -1;
}

void initialise_internal_node()
{
	int i=0;
	internal_node = (int *) calloc((num_samples+1),sizeof(int));
	for(i=0; i< num_samples; i++)
	{
		internal_node[i] = 0;
	}
}

void initialise_statistics()
{
	int i=0;
	statistics_for_samples = (sample_statistics **) calloc((num_samples+1),sizeof(sample_statistics *));
	for(i=0; i< num_samples; i++)
	{
		sample_statistics * sample_statistics_placeholder;
		sample_statistics_placeholder = (sample_statistics *) calloc(1, sizeof(sample_statistics));
		
		sample_statistics_placeholder->sample_name = (char *) calloc( MAX_SAMPLE_NAME_SIZE,sizeof(char));
		memcpy(sample_statistics_placeholder->sample_name, phylip_sample_names[i], size_of_string(phylip_sample_names[i]) +1);
		
		statistics_for_samples[i] = sample_statistics_placeholder;
	}
}

int number_of_snps_in_phylip()
{
	return num_snps;
}

void load_sequences_from_multifasta_file(char filename[])
{
    int i;

    num_snps    = genome_length(filename);
    num_samples = number_of_sequences_in_file(filename);

    sequences = (char **) calloc((num_samples+1),sizeof(char *));
    phylip_sample_names = (char **) calloc((num_samples+1),sizeof(char *));

    for(i = 0; i < num_samples; i++)
    {
        sequences[i] = (char *) calloc((num_snps+1),sizeof(char));
        phylip_sample_names[i] = (char *) calloc((MAX_SAMPLE_NAME_SIZE+1),sizeof(char));
    }
    get_sample_names_for_header(filename, phylip_sample_names, num_samples);

    int l;
    int sequence_number = 0;

    gzFile fp;
    kseq_t *seq;

    fp = gzopen(filename, "r");
    seq = kseq_init(fp);

    while ((l = kseq_read(seq)) >= 0)
    {
        for(i = 0; i< num_snps; i++)
        {
            sequences[sequence_number][i] = toupper(((char *) seq->seq.s)[i]);
            if(sequences[sequence_number][i] == 'N')
            {
                sequences[sequence_number][i]  = '-';
            }
        }
        sequence_number++;
    }

    kseq_destroy(seq);
    gzclose(fp);

    initialise_statistics();
    initialise_internal_node();
}

void freeup_memory()
{
    int i;
    for(i = 0; i < num_samples; i++)
    {
        free(sequences[i]);
        free(phylip_sample_names[i]);
    }
    free(sequences);
    free(phylip_sample_names);
    free(internal_node);
}

