package stanford;

import java.io.*;
import java.util.*;
import java.util.regex.*;

import edu.stanford.nlp.io.IOUtils;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.process.AbstractTokenizer;
import edu.stanford.nlp.process.CoreLabelTokenFactory;
import edu.stanford.nlp.process.PTBTokenizer;
import edu.stanford.nlp.ling.CoreAnnotations.TextAnnotation;

import edu.stanford.nlp.ie.machinereading.domains.bionlp.*;

public class StanfordEventTokenizer {
	// [input] should be the path to raw_text files
	// [file_index] should contain the corpus name along with the file id.
	public static void tokenize(String input, String output, String file_index) throws Exception{
		BufferedReader br = new BufferedReader(new FileReader(file_index));
		String fileName;
		while ((fileName= br.readLine()) != null){
//			filenames.add(line);
			String rawText = IOUtils.slurpFile(new File(input+fileName+".txt"));
			String output_file_name = output+fileName+".tok";
			File outputDir = new File(output_file_name.substring(0, output_file_name.lastIndexOf("/")));
			if (!outputDir.exists()){
				System.out.println("create directory:"+outputDir);
				outputDir.mkdirs();
			}
			BufferedWriter os = new BufferedWriter(new FileWriter(new File(output_file_name)));
			Collection<String> sentences = BioNLPFormatReader.splitSentences(rawText);
			for (String sentence: sentences){
				BioNLPTokenizer tokenizer = new BioNLPTokenizer(sentence);
				List<CoreLabel> tokens = tokenizer.tokenize();
				for (CoreLabel l: tokens){
					os.write(l.word()+" ");
				}
				os.newLine();
			}
			os.close();
		}
		br.close();
		
	}
//	public static void main(String[] args) throws Exception{
////		String rawTextPath = "../../data/BioNLP11/train/";
////		String rawTextPath = "../../data/BioNLP11/development/";
////		String rawTextPath = "../../data/BioNLP11/test/";
////		String outPath = "../BioNLPNN_Multi_Preprocess/cache/BioNLP11/parsed/stanford_extractor/tokenised/test/";
//		
////		String rawTextPath = "../../data/BioNLP11_Test/train/";
//		String rawTextPath = "../../data/BioNLP11_Test/test/";
//		String outPath = "../BioNLPNN_Multi_Preprocess/cache/BioNLP11_Test/parsed/stanford_extractor/tokenised/";
//		
////		String rawTextPath = "../../data/Open-access-pubmed-NFkB/";
////		String outPath = "../../data/Open-access-pubmed-NFkB_tokenised_by_stanford_extractor/";
////		String rawTextPath = "../../data/Open-access-pubmed-NFkB-2/";
////		String outPath = "../../data/Open-access-pubmed-NFkB-2_tokenised_by_stanford_extractor/";
//		
////		String rawTextPath = "../../data/BioNLP13/train/";
////		String rawTextPath = "../../data/BioNLP13/development/";
////		String outPath = "../BioNLPNN_Multi_Preprocess/cache/BioNLP13/parsed/stanford_extractor/tokenised/";
//		
////		String rawTextPath = "../../data/MergeBioNLP_Test/train/";
////		String rawTextPath = "../../data/MergeBioNLP_Test/test/";
////		String rawTextPath = "../../data/MergeBioNLP/development/maskPMID/";
////		String outPath = "../BioNLPNN_Multi_Preprocess/cache/MergeBioNLP/parsed/stanford_extractor/tokenised/";
////		String outPath = "../BioNLPNN_Multi_Preprocess/cache/MergeBioNLP_Test/parsed/stanford_extractor/tokenised/";
//		System.out.println("current dir:"+new File(".").getAbsolutePath());
//		System.out.println("parse files under "+rawTextPath);
//		File dir = new File(rawTextPath);
//		System.out.println("absolute path:"+dir.getAbsolutePath());
//		File[] files = dir.listFiles(new FilenameFilter(){
//			public boolean accept(File dir, String filename){return filename.endsWith(".txt");}
//		});
//		
////		BioNLPFormatReader reader = new BioNLPFormatReader();
//		File outputDir = new File(outPath);
//		if (!outputDir.exists()){
//			System.out.println("create directory:"+outputDir);
//			outputDir.mkdir();
//		}
//		for (File file: files){
//			String rawText = IOUtils.slurpFile(file);
////			rawText.replace("α", "alpha");
////			rawText.replace("β", "beta");
////			rawText.replace("γ", "gamma");
////			rawText.replace("-","-");
////			rawText.replace("ï","i");
////			rawText.replace("≤","<=");
////			rawText.replace("±","+/-");
////			rawText.replace("™","(TM)");
////			rawText.replace("×","x");
////			rawText.replace(" "," ");
////			rawText.replace("µ","mu");
////			rawText.replace("°","degrees");
////			rawText.replace("′","'");
//			
//			String fileName = file.getName();
//			System.out.println("process "+fileName);
//			fileName = outPath+fileName.substring(0, fileName.lastIndexOf(".txt"))+".tok";
//			BufferedWriter os = new BufferedWriter(new FileWriter(new File(fileName)));
//			Collection<String> sentences = BioNLPFormatReader.splitSentences(rawText);
//			for (String sentence: sentences){
//				BioNLPTokenizer tokenizer = new BioNLPTokenizer(sentence);
//				List<CoreLabel> tokens = tokenizer.tokenize();
//				for (CoreLabel l: tokens){
//					os.write(l.word()+" ");
//				}
//				os.newLine();
//			}
//			os.close();
////			System.out.println("write into "+fileName);
//		}
//		
//	}
	
	protected static class BioNLPTokenizer {
	    AbstractTokenizer<CoreLabel> tokenizer;
	    CoreLabelTokenFactory tokenFactory;
	    
	    /** If true, remove tokens that are standalone dashes. They are likely to hurt the parser. */
	    private static final boolean DISCARD_STANDALONE_DASHES = true;
	    
	    // XXX the following list is domain specific
	    
	    // this is a list of all suffixes which we'll split off if they follow a dash
	    // e.g. "X-induced" will be split into the three tokens "X", "-", and "induced"
			private static final HashSet<String> VALID_DASH_SUFFIXES = new HashSet<String>(
			    Arrays.asList(new String[] { "\\w+ed", "\\w+ing", "(in)?dependent",
			        "deficient", "response", "protein", "by", "specific", "like",
			        "inducible", "responsive", "gene", "mRNA", "transcription",
			        "sensitive", "bound", "driven", "positive", "negative", "dominant",
			        "family", "resistant", "activity", "proximal", "defective" }));
	    private final Pattern dashSuffixes;

		// matches a word followed by a slash followed by a word
	    private static final Pattern ANYSLASH_PATTERN = Pattern.compile("(\\w+)(/)(\\w+)");
	    
		public BioNLPTokenizer(String buffer) {
			StringReader sr = new StringReader(buffer);
			String options = "ptb3Escaping=false";
			tokenFactory = new CoreLabelTokenFactory();
			tokenizer = new PTBTokenizer<CoreLabel>(sr, tokenFactory, options);
	
			String allSuffixes = makeRegexOr(VALID_DASH_SUFFIXES);
			String allSuffixesRegex = "(\\w+)(-)(" + allSuffixes + ")";
			dashSuffixes = Pattern.compile(allSuffixesRegex, Pattern.CASE_INSENSITIVE);
		}

	    private String makeRegexOr(Iterable<String> pieces) {
	      StringBuilder suffixBuilder = new StringBuilder();
	      for (String suffix : pieces) {
	      	if (suffixBuilder.length() > 0) {
	      		suffixBuilder.append("|");
	      	}
	      	suffixBuilder.append("(" + suffix + ")");
	      }
	      return suffixBuilder.toString();
	    }

		public List<CoreLabel> tokenize() {
			List<CoreLabel> tokens = tokenizer.tokenize();
			for (CoreLabel token:tokens){
				if ((token.endPosition()-token.beginPosition())<token.word().length()){
					for (CoreLabel tok:tokens){
						System.out.print(tok.word()+" ");
					}
					System.out.println();
					System.out.print("tranform "+token.word() +" to ");
					token.setWord(token.word().substring(0,token.endPosition()-token.beginPosition()));
					System.out.println(token.word());
				}
			}
			return postprocess(tokens);
		}

	    protected List<CoreLabel> postprocess(List<CoreLabel> tokens) {
			// we can't use the regex "(anti)|(non)" since that will create an extra
			// group and confuse breakOnPattern, thus we do an extra pass
			tokens = breakOnPattern(tokens, Pattern.compile("(anti)(-)(\\w+)",
				Pattern.CASE_INSENSITIVE));
			tokens = breakOnPattern(tokens, Pattern.compile("(non)(-)(\\w+)",
					Pattern.CASE_INSENSITIVE));
			tokens = breakOnPattern(tokens, dashSuffixes);
//			tokens = breakOnPattern(tokens, ANYSLASH_PATTERN);
			tokens = breakOnPattern(tokens);//, ANYSLASH_PATTERN
			
			// re-join trailing or preceding - or + to previous digit
			tokens = joinSigns(tokens);
			
			// convert parens to normalized forms, e.g., -LRB-. better for parsing
			tokens = normalizeParens(tokens);
			
			return tokens;
	    }
	    
	    private List<CoreLabel> joinSigns(List<CoreLabel> tokens) {
	      List<CoreLabel> output = new ArrayList<CoreLabel>();
	      for (int i = 0; i < tokens.size(); i++) {
	        // -/-
	        if(i < tokens.size() - 3 &&
	            tokens.get(i).endPosition() == tokens.get(i + 1).beginPosition() && // neighbor
	            tokens.get(i + 1).word().equals("-") && // -/-
	            tokens.get(i + 2).word().equals("/") &&
	            tokens.get(i + 3).word().equals("-")){
	          String word = tokens.get(i).word() +
	            tokens.get(i + 1).word() +
	            tokens.get(i + 2).word() +
	            tokens.get(i + 3).word();
	          output.add(tokenFactory.makeToken(word, tokens.get(i).beginPosition(), word.length()));
	          i += 3;
	          continue;
	        }
	        
	        // - or +
	        if(i < tokens.size() - 1){
	          CoreLabel crt = tokens.get(i);
	          CoreLabel nxt = tokens.get(i + 1);
	          
	          // trailing +
	          if(crt.endPosition() == nxt.beginPosition() && // neighbor
	              ! isParen(crt.word()) && // not '('
	              nxt.word().equals("+")){ // '[TOKEN]+'
	            String word = crt.word() + nxt.word();
	            output.add(tokenFactory.makeToken(word, crt.beginPosition(), word.length()));
	            i ++;
	            continue;
	          }
	          
	          // trailing -
	          if(crt.endPosition() == nxt.beginPosition() &&
	              (i + 2 >= tokens.size() || nxt.endPosition() != tokens.get(i + 2).beginPosition()) && // not A-B, shoud be A-$ or A- B, to distinct the sign or hyphen
	              ! isParen(crt.word()) && // current is not '('
	              nxt.word().equals("-")){ // '[TOKEN]-'
	            String word = crt.word() + nxt.word();
	            output.add(tokenFactory.makeToken(word, crt.beginPosition(), word.length()));
	            i ++;
	            continue;
	          }
	          
	          // preceding -
	          if(crt.endPosition() == nxt.beginPosition() &&
	              (i == 0 || crt.beginPosition() != tokens.get(i - 1).endPosition()) &&
	              ! isParen(nxt.word()) &&
	              crt.word().equals("-")){
	            String word = crt.word() + nxt.word();
	            output.add(tokenFactory.makeToken(word, crt.beginPosition(), word.length()));
	            i ++;
	            continue;
	          }
	        }
	        
	        output.add(tokens.get(i));
	      }
	      return output;
	    }
	    
	    private static final String [][] PARENS = {
	      { "(", "-LRB-" },
	      { ")", "-RRB-" },
	      { "[", "-LSB-" },
	      { "]", "-RSB-" }
	    };
	    
	    private static boolean isParen(String s) {
	      for(int j = 0; j < PARENS.length; j ++){
	        if(s.equals(PARENS[j][0])){
	          return true;
	        }
	      }
	      return false;
	    }
	    
	    private static List<CoreLabel> normalizeParens(List<CoreLabel> tokens) {
	      for (int i = 0; i < tokens.size(); i++) {
	        CoreLabel token = tokens.get(i);
	        for(int j = 0; j < PARENS.length; j ++){
	          if(token.word().equals(PARENS[j][0])){
	            token.set(TextAnnotation.class, PARENS[j][1]);
	            }
	        }
	      }
	      return tokens;
	    }

		private List<CoreLabel> breakOnPattern(List<CoreLabel> tokens, Pattern pattern) {
			List<CoreLabel> output = new ArrayList<CoreLabel>();
			for (int i = 0; i < tokens.size(); i++) {
				CoreLabel token = tokens.get(i);
				Matcher matcher = pattern.matcher(token.word());
				if (matcher.find()) {
					int sepPos = matcher.start(2);
					String s1 = token.word().substring(0, sepPos);
					if(! DISCARD_STANDALONE_DASHES || ! s1.equals("-")){
					  output.add(tokenFactory.makeToken(s1, token.beginPosition(), sepPos));
					}
					String sep = matcher.group(2);
					if(! DISCARD_STANDALONE_DASHES || ! sep.equals("-")){
					  output.add(tokenFactory.makeToken(sep, token.beginPosition() + sepPos, 1));
					}
					String s3 = token.word().substring(sepPos + 1);
					if(! DISCARD_STANDALONE_DASHES || ! s3.equals("-")){
						output.add(tokenFactory.makeToken(s3, token.beginPosition() + sepPos + 1, 
						token.endPosition() - token.beginPosition() - sepPos - 1));
					}

				} else {
				  output.add(token);
				}
			}
			return output;
		}
	    private List<CoreLabel> breakOnPattern(List<CoreLabel> tokens) {
	    	Pattern pattern = ANYSLASH_PATTERN;
			List<CoreLabel> output = new ArrayList<CoreLabel>();
			for (int i = 0; i < tokens.size(); i++) {
				CoreLabel token = tokens.get(i);
				while(true){
					Matcher matcher = pattern.matcher(token.word());
					if (matcher.find()) {
						int sepPos = matcher.start(2);
						String s1 = token.word().substring(0, sepPos);
						if(! DISCARD_STANDALONE_DASHES || ! s1.equals("-")){
						  output.add(tokenFactory.makeToken(s1, token.beginPosition(), sepPos));
						}
						String sep = matcher.group(2);
						if(! DISCARD_STANDALONE_DASHES || ! sep.equals("-")){
						  output.add(tokenFactory.makeToken(sep, token.beginPosition() + sepPos, 1));
						}
						String s3 = token.word().substring(sepPos + 1);
						if(! DISCARD_STANDALONE_DASHES || ! s3.equals("-")){
							token = tokenFactory.makeToken(s3, token.beginPosition() + sepPos + 1, 
							token.endPosition() - token.beginPosition() - sepPos - 1);
						}
					} else {
						output.add(token);
						break;
					}
				}
			}
			return output;
		}
	  }
}
