#!/usr/local/bin/perl
# Rene Warren (c) 2008

use strict;
use Math::TrulyRandom;
use Data::Dumper;
use lib "/home/rwarren/perl/Production/SAM/lib";
use strict;
use SequenceAnalysis;


if($#ARGV<3){
   die "Usage $0 <V genes fasta> <J genes fasta> <JC sequence relationship> <Number of TCR to generate>\n";
}

my $jdel_mod = "Jdel.model";
my $vdel_mod = "Vdel.model";
my $dadd_mod = "Dadd.model"; 
my $dbase_mod = "Dbase.model";

if(! -e $jdel_mod){die "Can't find config file  $jdel_mod -- fatal.\n";}
if(! -e $vdel_mod){die "Can't find config file  $vdel_mod -- fatal.\n";}
if(! -e $dadd_mod){die "Can't find config file  $dadd_mod -- fatal.\n";}
if(! -e $dbase_mod){die "Can't find config file  $dbase_mod -- fatal.\n";}

my $v_file = $ARGV[0];
my $j_file = $ARGV[1];
my $jc_file = $ARGV[2];
my $tcr_number = $ARGV[3]; 

my $max = $tcr_number;
my $QCtcr = 0;

print "Reading-in gene sequences..\n";
my $v_arr = &readSequences($v_file);
my $j_arr = &readSequences($j_file);

########################## MODELS ########################
open (JDEL,$jdel_mod);
my $j_csv = <JDEL>;
chomp($j_csv);
my @j_del = split(/\,/,$j_csv);
my $jdelct = $#j_del;          
close JDEL;

open (VDEL,$vdel_mod);
my $v_csv = <VDEL>;
chomp($v_csv);
my @v_del = split(/\,/,$v_csv);
my $vdelct = $#v_del;
close VDEL;

open (DBASE,$dadd_mod);
my $d_csv = <DBASE>;
chomp($d_csv);
my @d_base = split(/\,/,$d_csv);
my $dbasect = $#d_base;
close DBASE;

open (D,$dbase_mod);
my $dct=0;
my $d_track;
while(<D>){
   chomp;
   my @d_add = split(/\,/);   
   $d_track->{$dct}=\@d_add;
   $dct++;
}
close D;
########################## MODELS ########################

open(JC,$jc_file);
my $jc_hash;
while(<JC>){
   chomp;
   my @a=split(/\:/);
   $jc_hash->{$a[0]} = $a[1];
}
close JC;

my @v_arr = @$v_arr;
my @j_arr = @$j_arr;


#### Initiate Sequence
print "Preparing Randomization of TCR.  Target is n=$tcr_number distinct sequences.\n";
srand(time() ^($$ + ($$ <<15)));
my $vnum = $#v_arr+1;
my $jnum = $#j_arr+1;
print "\tThere are $vnum V gene sequences and $jnum J gene sequences in your input files.\n"; 

#### Will do this as long as user want sequence
open(OUT, ">TCRtable1.fa") || die "Can't open file for writing -- fatal\n";

while ($QCtcr < $max){

   #print "Building sequence # $QCtcr\n";

   #### RANDOM SELECTION OF SEQUENCE SEGMENTS V AND J
   my $vrand = int(rand($vnum));
   #print "Getting V-gene # $vrand --->$v_arr[$vrand]<---\n";
   my $jrand = int(rand($jnum));
   #print "Getting J-gene # $jrand --->$j_arr[$jrand]<---\n";
 
   #### RANDOM SELECTION OF NUMBER OF BASES TO DELETE BASED ON FREQUENCY MODEL || V-3' and 5'-J
   my $vdelrand = int(rand($vdelct));
   #print "Deleting V-3' # $vdelrand --->$v_del[$vdelrand]<---\n";
   my $jdelrand = int(rand($jdelct));
   #print "Deleting 5'-J # $jdelrand --->$j_del[$jdelrand]<---\n";

   #### RANDOM SELECTION OF BASES TO ADD BETWEEN V-J / BASE SELECTED AT RANDOM BASED ON FREQUENCY MODEL
   my $dbaserand = int(rand($dbasect));
   #print "Adding # $dbaserand --->$d_base[$dbaserand]<---\n";

   my $pos = 0;
   my $Dseq = "";
   while($pos < $d_base[$dbaserand]){
      my $tmp = $d_track->{$pos};
      my @choice = @$tmp;
      my $choice_ct = $#choice;
      my $brand = int(rand($choice_ct));
      $Dseq .= $choice[$brand];
      #print "$pos. RANDOM ($brand) BASE: $choice[$brand]\n";
      $pos++;
   }   

   #### BUILDING SEQUENCE


   #### CHOOSING RIGHT C BASED ON J
   my $v_length = length($v_arr[$vrand]);
   my $v_tnt = $v_length + $v_del[$vdelrand];
   my $new_v = substr($v_arr[$vrand],0,$v_tnt);

   my $j_length = length($j_arr[$jrand]);
   my $j_tnt = $v_length + $j_del[$jdelrand];
   my $new_j = substr($j_arr[$jrand],abs($j_del[$jdelrand]),$j_tnt);

   #print "$v_arr[$vrand] $v_del[$vdelrand]\n";
   #print "$new_v\n\n";
   #print "$j_arr[$jrand] $j_del[$jdelrand]\n";
   #print "$new_j\n\n";
   #print "$v_del[$vdelrand] -($Dseq)- $j_del[$jdelrand]\n"; 

   my $TCR = $new_v . $Dseq . $new_j . $jc_hash->{$j_arr[$jrand]};

   my $TCR_size = length($TCR);
   if($TCR_size % 3 == 0){

       my $TCRaa=SequenceAnalysis::dna2protein($TCR);
       if($TCRaa =~ /\*/){
          #print "STOP CODON: $TCRaa\n";
       }else{        
          if($TCR =~ /([AC][ACGT]{3}[CAG][ACGT][GC][AC][ACGT][AT][CTG][ACTG][ATG][GC][ATC][ACTG][TAC][ACTG]TA\w+GCTGTGTTTGAGCCATCAGA)\w+/){  ### That's the PCR primer - since it is a 5' RACE, don't care about stuff in 3'
             my $tcr_amp = $1;

             if(length($tcr_amp)<=190){
                $QCtcr++; #### QC ON FORMED SEQUENCES (NO STOP CODON IN SEQ ALLOWED, LENGTH MULTIPLE OF 3)\n";
                          #### WILL PCR DETECT CRAP TCR, THAT DO NOT LEAD TO ANYTHING?

                #my $spos = length($tcr_amp) - 190;
                #my $size_select = substr($tcr_amp,$spos,190);
                #print OUT ">$QCtcr\n$size_select\n";

                print OUT ">$QCtcr:$v_del[$vdelrand]:$j_del[$jdelrand]:$d_base[$dbaserand]\n$tcr_amp\n";
             }
          }
       }
   #}else{
      #print "IMPROPER FRAME FOR $TCR ($TCR_size)\n";
   }

   #my $sleep_time=rand();
   #sleep($sleep_time);
}
close OUT;

#### Clonal Amplification of certain species, chosen at random
#### Generate reads with errors from corresponding concatenated amplicons


#-----------------------------------------------------
sub readSequences{
   my $fasta = shift;
   
   my @reads;
   open (FASTA, $fasta) or die "can't open $fasta\n";
   my $ct=0;
   my $prev = "";
   my $seq = "";
   while (<FASTA>){
      chomp;
      if (/\>(\S+)/){
         my $head = $1;
         if($head ne $prev && $seq ne "" && $prev ne ""){
            #$ct++;
            #print "$ct $seq\n";
            push @reads, $seq;
         }
         $seq = "";
         $prev = $head;
      }elsif(/([ACGT]*)/i){
         $seq .= $1;
      }
   }
   $ct++;
   push @reads, $seq;
   #print "$ct. $seq\n";

   close FASTA;

   return \@reads;
}

