#!/usr/bin/env perl
use JSON;
use HTTP::Tiny;
use File::Basename;

BEGIN {
  if ($0 =~ /([^(\/)]+)$/) {
    push (@INC, "$`lib/");
  }
}
require "RSA.lib";
use RSAT::server;

## Main package
package main;
{
  ################################################################
  ## Initialize parameters
  our $start_time = RSAT::util::StartScript();
  our $verbose = 0;

  our $ensembl_ftp_server = 'http://ftp.ebi.ac.uk';
  our $ensembl_ftp_path = 'pub/ensemblorganisms';
  our $ensembl_ftp_url = "$ensembl_ftp_server/$ensembl_ftp_path";
  our $ensembl_ftp_jsonfile = "$ensembl_ftp_url/species.json";
  our $ensembl_ftp_local_jsonfile = $ENV{RSAT}."/downloads/species.json";

  # local folder for genome data
  our $local_dir = $ENV{RSAT}."/public_html/data/genomes/";

  our $get_command = "wget";
  our $die_on_error = 1;

  ## Tasks to handle download and parsing of genomes from Ensembl
  our @tasks = qw(
      available
      download
      parse
      install
      default);
  our @default_tasks = qw(download parse install);
  our $supported_tasks = join(",", @tasks);

  our $parse_subtasks  = "split_features,parse_gtf,seq_len,parse_fasta";
  our $install_subtasks = "allup,upstream_freq,oligos,dyads,clean,config";

  my $null = "<NA>";
  my ($source,$taxonomy,$taxid) = ($null,$null,$null);
  our (%infile, %dir, %task, %downloads, %taxon, %found, @user_tasks);

  ## File containing a list of organisms to install from Ensembl
  $infile{organisms} = "";

  ReadArguments();

  ################################################################
  ## Check parameter values

  if(scalar(@user_tasks) == 0) {
    RSAT::error::FatalError("You should specify at least one task.\nSupported tasks\n\t$supported_tasks\n");

  } elsif(!$task{available} && (scalar(@organisms)==0 && !$infile{organisms})) {
    RSAT::error::FatalError("You should choose at least one organism to install (-org, -org_file)\n");

  } elsif ($task{'default'}) {
    foreach my $t (@default_tasks) {
      $task{$t} = 1;
    }
    RSAT::message::Info("Default tasks selected: ", join(",", @default_tasks)) if ($main::verbose >= 2);
  }

  RSAT::util::CheckOutDir($local_dir);
  
  our %remaining_task = %task;

  if ($infile{organisms}) {
      ## If an organism file has been specified, read organism list
      my ($orgs) = &OpenInputFile($infile{organisms});
      while (<$orgs>) {
          next if (/^;/ || /^#/); ## Skip comment/header
          if(/^(\S+)/){  
            push(@organisms, $1)
	  }
      }
  }

  ################################################################
  ## Carry out tasks

  my ($genomeID);

  if ($task{available} || $task{download} || $task{parse} || $task{install}) {

    my ($data,$response,$json_string);
    my ($species_name,$species,$taxid,$assembly_id,$assembly,$assembly_name);
    my ($strain,$provider,$provider_data,$release,$release_node,$paths);
    my ($hmasked,$unmasked,$genes_gtf,$pep,$vcf);
    my ($md5file,$hmasked_idx,$unmasked_idx,$pep_idx);

    # make sure local Ensembl species file is up to date
    $response = HTTP::Tiny->new->mirror($ensembl_ftp_jsonfile, $ensembl_ftp_local_jsonfile);

    open(JSON,'<',$ensembl_ftp_local_jsonfile) ||
      RSAT::error::FatalError("Failed downloading $ensembl_ftp_url_jsonfile\n");
    $/ = undef;
    $json_string = <JSON>;
    close(JSON);
    $/="\n";

    $data = decode_json($json_string);
    $json_string = '';

    # loop across current species and extract relevants bits from JSON
    foreach $species_name (sort keys %{ $data->{species} }) {

      $species = $data->{species}{$species_name};
      $taxid = $species->{taxid} // '';
      $strain = $species->{strain} // '';

      foreach $assembly_id (sort keys %{ $species->{assemblies} }) {

        $assembly = $species->{assemblies}{$assembly_id};
        $assembly_name = $assembly->{name} // '';

        foreach $provider (sort keys %{ $assembly->{genebuild_providers} }) {
          $provider_data = $assembly->{genebuild_providers}{$provider};

          foreach $release (sort keys %{$provider_data}) {
            $release_node = $provider_data->{$release};
            $genomeID = "$species_name.$assembly_id.$release.$provider";

	    if($task{'download'} || $task{parse} || $task{install}) {
		   
              if(grep(/^$genomeID$/, @organisms)) {

                RSAT::message::Info("Found genomeID $genomeID") if ($main::verbose >= 2);
                $found{$genomeID} = 1;		

                $paths = $release_node->{paths};

		# bgzipped files, to be renamed as .gz
                $hmasked =
                  $assembly->{assembly}{files}{genome_sequences}{'hardmasked.fa.bgz'} // '';
                $unmasked = 
                  $assembly->{assembly}{files}{genome_sequences}{'unmasked.fa.bgz'} // '';
                $genes_gtf =  
	          $paths->{genebuild}{files}{annotations}{'genes.gtf.bgz'} // '';
                $pep =
                  $paths->{genebuild}{files}{annotations}{'pep.fa.bgz'} // '';
                $vcf =
                  $paths->{short_variants}{files}{variation_data}{'variation.vcf.gz'} // '';

		# files not listed in JSON, but found in FTP site
		$md5file      = $assembly->{assembly}{files}{genome_sequences}{'hardmasked.fa.bgz'};
		$md5file      =~ s/hardmasked.fa.bgz/md5sum.txt/; 
		$hmasked_idx  = $assembly->{assembly}{files}{genome_sequences}{'hardmasked.fa.bgz'} . '.fai';
		$unmasked_idx = $assembly->{assembly}{files}{genome_sequences}{'unmasked.fa.bgz'} . '.fai';
		$pep_idx      = $paths->{genebuild}{files}{annotations}{'pep.fa.bgz'} . '.fai';

                $downloads{$genomeID} = [
		  $md5file,
		  $unmasked,$hmasked,$pep,$genes_gtf,
		  $vcf,
                  $hmasked_idx,$unmasked_idx,$pep_idx
	        ];

		$filenames{$genomeID} = [
                  'md5sum.txt',			
                  $genomeID.'.dna.genome.fa.gz',
                  $genomeID.'.dna_rm.genome.fa.gz',
                  $genomeID.'.pep.all.fa.gz',
                  $genomeID.'.gtf.gz',
		  $genomeID.'.vcf.gz',
		  $genomeID.'.dna_rm.genome.fa.gz.fai',
                  $genomeID.'.dna.genome.fa.gz.fai',
                  $genomeID.'.pep.all.fa.gz.fai' 
	        ];

	        $taxon{$genomeID} = $taxid; 

              }	
            } else {
              print "$genomeID\t[$assembly_name, $strain]\n"
            }
          }
        }
      }
    }
  } 

  # any other remaining tasks
  my ($org_dir, $var_dir, $cmd, $cmd2, $f, $file, $rfile, $chrn);
  my ($localfile, $md5sum, %checksums, %chrnum);

  foreach $genomeID (@organisms) {

    if(!$found{$genomeID}) {
      print "# $genomeID not found, please check/correct genomeID\n"; 	    
      next;
    }

    if ($task{download}) {

      # create/check destination folder
      $org_dir = "$local_dir/$genomeID/genome/";
      RSAT::util::CheckOutDir($org_dir);
      RSAT::message::Info("downloading $genomeID to $org_dir") if ($main::verbose >= 2);

      if ($get_command eq "wget") {
        if ($main::verbose < 1) {
          $cmd = "$get_command -c --quiet";
        } else {
          $cmd = "$get_command -c -nv";
        } 

        foreach $f (0 .. $#{$downloads{$genomeID}}) {

          next if($downloads{$genomeID}->[$f] eq '');

          $file = $downloads{$genomeID}->[$f];
	  $localfile = $filenames{$genomeID}->[$f]; 

          if($localfile =~ m/\.vcf\.gz$/) { 
 
            # create variations/ only if needed
            $var_dir = "$local_dir/$genomeID/variations/";
            RSAT::util::CheckOutDir($var_dir); 

            # VCF file is stored in a var_dir
            doit("$cmd $ensembl_ftp_url/$file -O $var_dir$localfile", 0, $die_on_error, $main::verbose);
            sleep(1);

	    if(!-e "$var_dir$genomeID\.varBed") { 
	      $cmd2 = "convert-variations -i $var_dir$localfile -from vcf -to varBed -o $var_dir$genomeID\.varBed";
	      doit($cmd2, 0, $die_on_error, $main::verbose);

	      open(SOURCE,">$var_dir/source.txt") ||
                RSAT::error::FatalError("Cannot write $var_dir/source.txt\n");
              print SOURCE "Ensembl\n";
              close(SOURCE);	      
	    }

          } else { 
		  
            # other files in org_dir		  
            doit("$cmd $ensembl_ftp_url/$file -O $org_dir$localfile", 0, $die_on_error, $main::verbose);
            sleep(1);

	    if($localfile eq 'md5sum.txt') {

              my ($md5h) = RSAT::util::OpenInputFile($org_dir.$localfile);
              while(my $line = <$md5h>) { 
                if($line =~ m/^(\S+)\s+(\S+)/) {		      
                  ($md5sum,$rfile) = ($1,$2);
		  $rfile = basename($rfile);
                  $checksums{ $rfile } = $md5sum;
	        }
              }
	      close($md5h);
  
            } elsif($file =~ m/masked\.fa\.bgz$/) { 
	
	      # checksums and chr number of genome sequence files
	      $md5sum = (split(/\s+/,`md5sum $org_dir$localfile`))[0];
	      $chrn = `zgrep -c "^>" $org_dir$localfile`; chomp $chrn;
	      
	      if(!$chrnum{$genomeID}) {
                $chrnum{$genomeID} = $chrn

              } else {
                if($chrn != $chrnum{$genomeID}) {
                  print "# $genomeID task: download failed (chr number)\n\n"; 
                die "# $ensembl_ftp_url/$file\n" .
                  "# expected chr number: $chrnum{$genomeID}\n# download chr number: $chrn\n";
                }
              }
	      
	      if($md5sum ne $checksums{basename($file)}) {
                print "# $genomeID task: download failed (checksum)\n\n";
                die "# $ensembl_ftp_url/$file\n" .
		  "# declared checksum: ".$checksums{basename($file)}. "\n# download checksum: $md5sum\n";
              } 
            }
          }  
        }
      }

      print "# $genomeID task: download OK\n"; 
    }  

    if ($task{parse}) {
      $org_dir = "$local_dir/$genomeID/genome/";
      RSAT::util::CheckOutDir($org_dir);
      RSAT::message::Info("parsing $genomeID files in $org_dir") if ($main::verbose >= 2);

      $cmd = "parse-gtf " .
        "-i $org_dir$genomeID.gtf.gz " .
        "-fasta_rm $org_dir$genomeID.dna_rm.genome.fa.gz " . 
        "-fasta $org_dir$genomeID.dna.genome.fa.gz " .
        "-fasta_pep $org_dir$genomeID.pep.all.fa.gz " .
        "-org_name $genomeID " .
        "-taxid $taxon{$genomeID} " .
        "-gtf_source Ensembl " .
        "-task $parse_subtasks " .	
        "-v $main::verbose " .
        "-o $org_dir ";
     
      doit($cmd, 0, $die_on_error, $main::verbose);
      RSAT::message::TimeWarn($cmd) if ($main::verbose >= 3);
    }
      
    if ($task{install}) {
      RSAT::message::Info("installing $genomeID") if ($main::verbose >= 2); 

      $cmd = "install-organism " .
        "-org $genomeID " .
        "-task $install_subtasks " .
        "-source Ensembl " .
	"-rm " .
        "-v $main::verbose ";
      
      doit($cmd, 0, $die_on_error, $main::verbose);
       RSAT::message::TimeWarn($cmd) if ($main::verbose >= 3);
    }
  }

  Finish($start_time);
}


################################################################
##################  SUBROUTINE DEFINITIONS  ####################
################################################################

sub Finish {
  my $start_time = (@_);	
  my $exec_time = RSAT::util::ReportExecutionTime($start_time);
  print $exec_time if ($main::verbose >= 1);
  exit(0);
}

################################################################
## Display full help message
sub PrintHelp {
    open HELP, "| more";
    print HELP <<End_of_help;
NAME
	install-organism-ensembl

AUTHOR
        Jacques van Helden (Jacques.van-Helden\@univ-amu.fr)
	Bruno Contreras Moreira (bcontreras\@eead.csic.es)

REQUIREMENTS

This program relies on several system tools: wget, md5sum, grep

USAGE
     
        install-organism-ensembl -org Arabidopsis_thaliana.GCA_000001735.1.2010_09.community_araport11

	install-organism-ensembl -task available | grep Arabidopsis_thaliana > mylist.txt
	install-organism-ensembl -org_file mylist.txt -task default

DESCRIPTION
	Install organisms from Ensembl on the local RSAT instance.

	This script is a task manager, which (depending on the
	selected tasks) manages different steps necessary for the
	installation of an organism from flat files:

	- Downloading genome files from Ensembl, including vcf files when available.
          ($ensembl_ftp_url).

	- Parsing files and calculating oligonucleotide and dyad frequencies, which will
          be used as background models for motif discovery.

	- Adding the organism to local config file.

	- Installing the organism.

CATEGORY
	Data management.

OPTIONS
	-h	(must be first argument) display full help message
	-help	(must be first argument) display options
	-v	verbose

    ARGUMENTS

        -org    genomeID as provided by -task available, without spaces.
                The option -org can be repeated on the same command line 
                to iterate the installation over multiple organisms. Example:
                
		-org Arabidopsis_thaliana.GCA_000001735.1.2010_09.community_araport11

                which corresponds to 'Arabidopsis_thaliana',
		ENA/NCBI assembly 'GCA_000001735.1' and
		annotated geneset '2010_09', provided by 'community_araport11'
                
        -org_file organism_file_name
                Text file containing a list of genomeIDs to install.
                Only the first word of each row is considered.

	-dir    Absolute path of the installation directory, 
                default: $local_dir 

	-task	specification of a single installation task e.g.
			
                install-organism-ensembl -task available

		supported tasks: $supported_tasks

		Description of the tasks
		------------------------

		available    list available Ensembl species, parses large JSON file;
		             this task cannot be combined with any others.

		download     download genome from the Ensembl FTP site with wget
		             ($ensembl_ftp_url)

		parse        parse genome data (sequence, features) with parse-gtf

		install	     compute upstream freqs and install organism

		default      same as download,parse,install

SEE ALSO

    download-organisms

	The program I<download-organism> copies RSAT-formatted
	genomes from a RSAT server or repository. 

 	If a genome is available on the RSAT server, it is recommended
	to use download-genomes in order to obtain it immediately in
	the RSAT format, rather than install-genomes.

End_of_help
    close HELP;
    exit(0);
}

################################################################
## Display short help message
sub PrintOptions {
  open HELP, "| more";
  print HELP <<End_short_help;
install_organism options
------------------------
-h		(must be first argument) display full help message
-help		(must be first argument) display options
-v		verbose
-org		1-word genomeID, can be repeated. Obtained with -task available
-org_file	organism_file_name
-dir		absolute path of the installation directory
-task		installation task ($supported_tasks)

End_short_help
  close HELP;
  exit;
}

################################################################
## Read arguments 
sub ReadArguments {
  my @arguments = @ARGV; ## create a copy to shift, because we need ARGV to report command line in &Verbose()
  while (scalar(@arguments) >= 1) {
    $arg = shift (@arguments);

    #    foreach my $a (0..$#ARGV) {
    ### verbose ###
    if ($arg eq "-v") {
      if (IsNatural($arguments[0])) {
	$main::verbose = shift(@arguments);
      } else {
	$main::verbose = 1;
      }

      ## detailed help
    } elsif ($arg eq "-h") {
      PrintHelp();

      ## list of options
    } elsif ($arg eq "-help") {
      PrintOptions();

      ## 1-word organism name/genomeID
    } elsif ($arg eq "-org") {
      my $organism_short_name = shift(@arguments);
      RSAT::error::FatalError( $organism_short_name, "Invalid organism name: cannot contain spaces") if ($organism_short_name =~ /\s/);
      RSAT::error::FatalError( $organism_short_name, "Invalid organism name: should not be empty.") unless ($organism_short_name =~ /\S/);
      RSAT::error::FatalError( $organism_short_name, "Invalid organism name: should contain at least one underscore.") unless ($organism_short_name =~ /_/);
      push(@organisms, $organism_short_name);

      ## Input file with organism list
    } elsif ($arg eq "-org_file") {
      $infile{organisms} = shift(@arguments);

      ## installation dir
    } elsif ($arg eq "-dir") {
      $local_dir = shift(@arguments);

      ## task selection
    } elsif ($arg eq "-task") {
	my @requested_tasks = split ",", shift(@arguments);
	foreach my $t (@requested_tasks) {
	    next unless $t;
	    if ($main::supported_tasks =~ m/$t/) {
		$main::task{$t} = 1;
		push(@main::user_tasks, $t);
	    } else {
		RSAT::error::FatalError("Unsupported task '$t'; supported: $main::supported_tasks");
	    }
      }

=pod

=cut

    } 
  }
}