#!/usr/bin/perl

# Gibson Assembly fragment design based on Picky probe output
# Copyright 2014-2015 Hui-Hsien Chou

# 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.
# http://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html

use strict;
use warnings;

my $usage = <<"USAGE";
usage: $0 opt_option min_len max_len sequences.picky sequences.seq
  opt_option: optimize the count or cost of the fragments
  min_len: shortest individual fragment length allowed
  max_len: longest individual fragment length allowed
  sequences.picky: Picky output that defined the unique junctions
  sequences.seq: original gene sequences that are to be broken up
USAGE

die $usage unless @ARGV == 5;

my $optimize_option = shift @ARGV;
my $min_gblock_size = shift @ARGV;
my $max_gblock_size = shift @ARGV;
my $left_linker =  "ACTCAAGCTT";  # HindIII site
my $right_linker = "GAATTCGCCC";  # EcoRI site

die "optimize options must be either count or cost"
    unless $optimize_option eq "cost" or $optimize_option eq "count";

my %table;
open PICKY, $ARGV[0];
while (<PICKY>) {
  chomp;
  my @fields = split/\t/;
  my ($probe, $target, $nontarget) = @fields[0, 2, 3];
  for (my $i=0; $i<$target; $i++) {
    $_ = <PICKY>;
    chomp;
    @fields = split/\t/;
    die "incorrect target info line $_" unless @fields == 7;
    die "incorrect target line format $_" unless $fields[0] eq 'U' or $fields[0] eq 'S';
    $table{$fields[6]} = [ ] unless exists $table{$fields[6]};
    push @{$table{$fields[6]}}, [ @fields[4, 5], $probe ];
  }
  for (my $i=0; $i<$nontarget; $i++) {
    $_ = <PICKY>;
    chomp;
    @fields = split/\t/;
    die "incorrect nontarget info line $_" unless @fields == 7;
    die "incorrect nontarget line format $_" unless $fields[0] eq '>' or $fields[0] eq '<';
  }
  $_ = <PICKY>;
  chomp;
  die "extra lines found $_" if length;
}
close PICKY;

my ($name, $seq);
my $total_cost = 0.0;
my %front;
my @junctions;

sub find_best_gblock_parts {
  my $j = shift @_;
  return (0.0, 0) unless $j>0;
  my $best_leng = 1000000;
  my $best_cost = 1000000.0;
  my @best_parts;
  for (@{$front{$j}}) {
    my ($cost, @parts) = find_best_gblock_parts($_);
    my $len = $junctions[$j][1]-$junctions[$_][0]+1;
    if ($len <= 60) {
      $cost += 0.2 * $len;
    } else {
      $cost += 0.3 * $len;
    }
    if ($optimize_option eq "cost" && $cost<$best_cost || 
	$optimize_option eq "count" && @parts<$best_leng) {
      $best_leng = @parts;
      $best_cost = $cost;
      @best_parts = @parts;
    }
  }
  return ($best_cost, @best_parts, $j);
}

sub breaking_gene {
  die "unscreened gene $name" unless exists $table{$name};
  die "already screened gene $name" unless $table{$name};
  @junctions = sort { $a->[0] <=> $b->[0] } @{$table{$name}};
  $seq =~ tr/acgtu/ACGTU/;
  for (@junctions) {
    die "inconsistent probe/sequence match" unless
	substr($seq, $_->[0], $_->[1]-$_->[0]+1) eq $_->[2];
    delete $_->[2];
  }
  map { map { $_ += length $left_linker } @{$_} } @junctions;
  $seq = $left_linker . $seq . $right_linker;
  my ($i, $j) = ( 0, 0);
  while ($seq =~ /(G{6,})/g) {
    print STDERR "$name has the following poly-G fragments longer than 5:\n" unless $i;
    print STDERR "$1\n";
    $i = 1;
  }
  @junctions = ( [ 0, length($left_linker)-1], @junctions,
		 [ length($seq)-length($right_linker), length($seq)-1] );
  $i = 0;
  for (@junctions) {
    print STDERR $i++, ": @{$_}\n";
  }
  %front = ( 0 => -1 );
  for ($i=0; $i<@junctions; $i++) {
    next unless exists $front{$i};
    for ($j=$i+1; $j<@junctions && $junctions[$j][1]-$junctions[$i][0]+1<=$max_gblock_size; $j++) {
      unless ($junctions[$j][1]-$junctions[$i][0]+1<$min_gblock_size) {
	$front{$j} = [ ] unless exists $front{$j};
	push @{$front{$j}}, $i;
	print STDERR "$j -> $i  ";
      }
    }
    print STDERR "\n";
  }
  unless (exists $front{@junctions-1}) {
    print STDERR "$name cannot be broken into gBlocks with $min_gblock_size to $max_gblock_size size fragments\n";
  } else {
    my ($best_cost, @parts) = find_best_gblock_parts($#junctions);
    my $sum_cost = 0.0;
    for ($i=0; $i<@parts-1; $i++) {
      my $gblock = substr($seq, $junctions[$parts[$i]][0],
			  $junctions[$parts[$i+1]][1]-$junctions[$parts[$i]][0]+1);
      my $cost;
      if (length $gblock<=60) {
	$cost = 0.2 * length $gblock;
      } else {
	$cost = 0.3 * length $gblock;
      }
      print ">${name}_${i}_\$$cost\n$gblock\n";
      $sum_cost += $cost;
    }
    die "sum cost $sum_cost does not match best_cost $best_cost" unless $sum_cost == $best_cost;
    $total_cost += $sum_cost;
  }
  $table{$name} = undef;
}

open GENES, $ARGV[1];
while (<GENES>) {
  chomp;
  if (/^>(\S+)/) {
    &breaking_gene if $name;
    ($name, $seq) = ($1, "");
  } else {
    $seq .= $_;
  }
}
close GENES;
&breaking_gene if $name;

print STDERR "Total cost will be $total_cost\n";
