package prok;

import java.util.ArrayList;
import java.util.HashMap;

import fileIO.FileFormat;
import fileIO.TextStreamWriter;
import server.ServerTools;
import shared.Tools;
import template.ThreadWaiter;

/** Crawls ncbi's ftp site to download genomes and annotations */
public class FetchProks {
	
	public static void main(String[] args){
		//ftp://ftp.ncbi.nih.gov:21/genomes/refseq/bacteria/
		
		String baseAddress=args[0];
		String out=args.length>1 ? args[1] : "stdout";
		if(args.length>2){
			maxSpeciesPerGenus=Integer.parseInt(args[2]);
			System.err.println("Set maxSpeciesPerGenus="+maxSpeciesPerGenus);
		}
		if(args.length>3){
			findBest=Tools.parseBoolean(args[2]);
			System.err.println("Set findBest="+findBest);
		}
		TextStreamWriter tsw=new TextStreamWriter(out, true, false, false, FileFormat.TEXT);
		tsw.start();

//		iterateOuter(baseAddress, tsw);
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
		
		int threads=7;
		ArrayList<ProcessThread> alpt=new ArrayList<ProcessThread>(threads);
		for(int i=0; i<threads; i++){
			alpt.add(new ProcessThread(contents, tsw, i, threads));
		}
		for(ProcessThread pt : alpt){pt.start();}
		boolean success=ThreadWaiter.waitForThreads(alpt);
		
		for(ProcessThread pt : alpt){
			totalSpecies+=pt.totalSpeciesT;
			totalGenus+=pt.totalGenusT;
			totalGenomes+=pt.totalGenomesT;
		}
		System.err.println("Total Genomes: "+totalGenomes);
		System.err.println("Total Species: "+totalSpecies);
		System.err.println("Total Genuses: "+totalGenus);
		
		tsw.poisonAndWait();
		assert(success);
	}
	
	static class ProcessThread extends Thread {
		
		ProcessThread(ArrayList<String> speciesList_, TextStreamWriter tsw_, int tid_, int threads_){
			speciesList=speciesList_;
			tsw=tsw_;
			tid=tid_;
			threads=threads_;
		}
		
		@Override
		public void run(){
			for(String s : speciesList){
//				if((s.hashCode()&Integer.MAX_VALUE)%threads==tid) {
//					processSpecies(s);
//				}
				
				//This way one thread handles an entire genus
				if((getGenus(s).hashCode()&Integer.MAX_VALUE)%threads==tid) {
					processSpecies(s);
				}
			}
		}
		
		void processSpecies(String species){
			String genus=getGenus(species);
			if(genus!=null){
				final int count=seen(genus, seen);
				
				if(maxSpeciesPerGenus<1 || count<maxSpeciesPerGenus){
					int found=examineSpecies(species, tsw);
					if(found>=1){
						totalSpeciesT++;
						totalGenomesT+=found;
						if(count==0){totalGenusT++;}
						put(genus, found, seen);
					}
				}else{
					if(verbose){System.err.println("same genus: "+species+"\n"+genus);}
				}
			}else{
				if(verbose){System.err.println("bad species: "+species+"\n"+genus);}
			}
		}
		
		final ArrayList<String> speciesList;
		final int tid;
		final int threads;
		//This is OK now that threads work on a per-genus basis
		HashMap<String, Integer> seen=new HashMap<String, Integer>();
		final TextStreamWriter tsw;
		
		int totalSpeciesT=0;
		int totalGenusT=0;
		int totalGenomesT=0;
	}
	
	static String getGenus(String path){
		//Candidatus_Hamiltonella
		String name=path.substring(path.lastIndexOf('/')+1);
		if(name.startsWith("Candidatus_")){name=name.substring("Candidatus_".length());}
		int under=name.indexOf('_');
		if(under>0){
			return name.substring(0, under);
		}else{
			return null;
		}
	}
	
	static String getSpecies(String path){
		//Candidatus_Hamiltonella
		String name=path.substring(path.lastIndexOf('/')+1);
		if(name.startsWith("Candidatus_")){name=name.substring("Candidatus_".length());}
		return name;
	}
	
	static int examineSpecies(String baseAddress, TextStreamWriter tsw){
		if(verbose){System.err.println("examineSpecies: "+baseAddress);}
		String speciesName=getSpecies(baseAddress);
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
//		System.err.println("B: "+contents);
		int found=0;
		for(String s : contents){
//			System.err.println(s);
			if(s.contains("reference")){
//				System.err.println("Looking at '"+s+"'");
				found+=examineAssemblies(s, tsw, speciesName);
			}
		}
		if(found>0){return found;}
		for(String s : contents){
//			System.err.println(s);
			 if(s.contains("latest_assembly_versions")){
//				System.err.println("Looking at '"+s+"'");
				 found+=examineAssemblies(s, tsw, speciesName);
			}
		}
		if(found>0){return found;}
		for(String s : contents){
//			System.err.println(s);
			if(s.contains("all_assembly_versions")){
//				System.err.println("Looking at '"+s+"'");
				found+=examineAssemblies(s, tsw, speciesName);
			}
		}
		return found;
	}
	
	static int examineAssemblies(String baseAddress, TextStreamWriter tsw, String speciesName){
		if(verbose){System.err.println("examineAssemblies: "+baseAddress);}
		if(findBest){
			String best=findBestAssembly(baseAddress);
			if(best!=null){
				int x=examineAssembly(best, tsw, speciesName);
				if(x>0){return x;}
			}
		}
		
		
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
//		System.err.println("C: "+contents);
		
		int found=0;
		for(String s : contents){
//			System.err.println(s);
			found+=examineAssembly(s, tsw, speciesName);
			if(found>0){break;}
		}
		return found;
	}
	
	/** Tries to find the assembly with the longest contig */
	static String findBestAssembly(String baseAddress){
		if(verbose){System.err.println("findBestAssembly: "+baseAddress);}
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
//		System.err.println("C: "+contents);
		long[] bestStats=null;
		String best=null;
		for(String s : contents){
//			System.err.println(s);
			long[] stats=calcStats(s);
			if(stats!=null){
				if(best==null || isBetter(bestStats, stats)){
					best=s;
					bestStats=stats;
				}
			}else if(best==null){
				best=s;
				bestStats=stats;
			}
		}
		return best;
	}
	
	//contigs, size, maxContig
	static boolean isBetter(long[] a, long[] b) {
		if(a==null){return true;}
		if(b[1]>2*a[1]){return true;}
		if(b[1]<2*a[1]){return false;}

		if(b[2]>a[2]){return true;}
		if(b[2]<a[2]){return false;}
		
		return b[0]<a[0];
	}
	
	static long[] calcStats(String baseAddress){
		if(verbose){System.err.println("calcStats: "+baseAddress);}
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
		String report=null;
		for(String s : contents){
			if(s.endsWith("_assembly_report.txt")){
				report=s;
				break;
			}
		}
		if(report==null){
			if(verbose){System.err.println("Could not find report for "+baseAddress);}
			return null;
		}
		if(verbose){System.err.println("Report: "+report);}
		ArrayList<String> data=null;
		for(int i=0; i<=retries && data==null; i++){
			try {
				data = ServerTools.readFTPFile(report);
			} catch (Exception e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
				try {
					Thread.sleep(Tools.mid(10000, i*1000, 1000));
				} catch (InterruptedException e1) {
					// TODO Auto-generated catch block
					e1.printStackTrace();
				}
			}
		}
		if(data==null){return null;}
		int contigs=0;
		long size=0;
		long max=0;
		for(String s : data){
			if(s!=null && s.length()>0){
				if(s.charAt(0)=='#'){
					//skip
				}else{
					String[] split=s.split("\t");
					contigs++;
					long len;
					try {
						len=Long.parseLong(split[8]);
					} catch (NumberFormatException e) {
						len=1;
					}
					size+=len;
					max=Tools.max(max, len);
				}
			}
		}
		return new long[] {contigs, size, max};
	}
	
	static int examineAssembly(String baseAddress, TextStreamWriter tsw, String speciesName){
		if(verbose){System.err.println("examineAssembly: "+baseAddress);}
		ArrayList<String> contents=ServerTools.listDirectory(baseAddress, retries);
//		System.err.println("D: "+contents);
		String gff=null;
		String fna=null;
		for(String s : contents){
//			System.err.println(s);
			if(!s.contains("_from_genomic")){
				if(s.endsWith("genomic.fna.gz")){fna=s;}
				else if(s.endsWith("genomic.gff.gz")){gff=s;}
			}
		}
		if(fna!=null && gff!=null){
			System.err.println("Printing: "+fna);
//			System.err.println(fna);
			synchronized(tsw){
				if(renameFiles){
					tsw.println("wget -q -O - "+fna+" > "+speciesName+".fna.gz");
					tsw.println("wget -q -O - "+gff+" > "+speciesName+".gff.gz");
				}else{
					tsw.println("wget -q "+fna);
					tsw.println("wget -q "+gff);
				}
				tsw.println();
			}
			return 1;
		}
		return 0;
	}
	
	static String makeSubAddress(String baseAddress, String extension){
		if(!baseAddress.endsWith("/")){baseAddress=baseAddress+"/";}
		String subAddress=baseAddress+extension.substring(extension.indexOf('/')+1);
		return subAddress;
	}
	
	static int seen(String s, HashMap<String, Integer> map){
//		synchronized(map){
			Integer x=map.get(s);
			return x==null ? 0 : x.intValue();
//		}
	}
	static void put(String s, int found, HashMap<String, Integer> map){
//		synchronized(map){
			int present=seen(s, map);
			map.put(s, present+found);
//		}
	}
	
	static boolean verbose=true;
//	static boolean allowSameGenus=false;
	static int maxSpeciesPerGenus=1;
	static boolean renameFiles=true;
	static int retries=40;
	static boolean findBest=false;
	
//	private static HashMap<String, Integer> seen=new HashMap<String, Integer>();
	
	static int totalSpecies=0;
	static int totalGenus=0;
	static int totalGenomes=0;

	private static final Integer one=1;
	
}
