package wektoryzacja;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import pl.wroc.pwr.file.ImageFolderScanner;
import pl.wroc.pwr.image.IGSImage;
import pl.wroc.pwr.image.IRGBImage;
import pl.wroc.pwr.image.converter.RGBtoGSConverter;
import pl.wroc.pwr.image.io.GenericLoader;
import pl.wroc.pwr.image.io.GenericSaver;
import pl.wroc.pwr.image.model.RGBImage;
import pl.wroc.pwr.imagechannel.IImageChannel;
import pl.wroc.pwr.imagechannel.basic.CloneChannel;
import pl.wroc.pwr.imagechannel.basic.Invert;
import pl.wroc.pwr.imagechannel.basic.MultiplyMaskValue;
import pl.wroc.pwr.imagechannel.draw.DrawLine;
import pl.wroc.pwr.imagechannel.draw.DrawX;
import pl.wroc.pwr.imagechannel.model.ImageChannel;
import pl.wroc.pwr.imagechannel.roi.AddROIToChannel;
import pl.wroc.pwr.imagechannel.segmentation.SplitSeparate;
import pl.wroc.pwr.imagechannel.skeletonization.BranchDetector;
import pl.wroc.pwr.imagechannel.skeletonization.FastParallelThinning;
import pl.wroc.pwr.imagechannel.skeletonization.TerminationDetector;
import pl.wroc.pwr.imagechannel.statistic.ImageChannelMean;
import pl.wroc.pwr.imagechannel.thresholding.BasicThreshold;
import pl.wroc.pwr.roi.IROI;
import pl.wroc.pwr.roi.basic.ROIFromThreshold;
import pl.wroc.pwr.roi.model.ROIPixels;
import pl.wroc.pwr.vector.distance.EuclideanDistance;

public class WektoryzacjaObrazu {
	
	private WektoryzacjaKrzywej wektoryzacjaKrzywej = new WektoryzacjaKrzywej(5);
	private MultiplyMaskValue maskowanie = new MultiplyMaskValue(0.9f);
	private AddROIToChannel rysowanieROI = new AddROIToChannel(1);
	private AddROIToChannel kasowanieROI = new AddROIToChannel(0);
	private ZapisGrafuWektoryzacji zapisGrafu = new ZapisGrafuWektoryzacji();
	private EuclideanDistance odleglosc = new EuclideanDistance();
		
	private int[] s1x = new int[]{-1, -1, -1, +1, +1, +1,  0,  0,  0,  0,  0,  0,   +1, +1, +1, +1, +1, +1, -1, -1, -1, -1, -1, -1};
	private int[] s1y = new int[]{ 0,  0,  0,  0,  0,  0, -1, -1, -1, +1, +1, +1,   +1, +1, +1, -1, -1, -1, +1, +1, +1, -1, -1, -1};
	private int[] s2x = new int[]{-2, -2, -2, +2, +2, +2,  0, -1, +1,  0, -1, +1,   +2, +2, +1, +2, +2, +1, -2, -2, -1, -2, -2, -1};
	private int[] s2y = new int[]{ 0, -1, +1,  0, -1, +1, -2, -2, -2, +2, +2, +2,   +2, +1, +2, -2, -1, -2, +2, +1, +2, -2, -1, -2};
	private int[] cst = new int[]{20, 24, 24, 20, 24, 24, 20, 24, 24, 20, 24, 24,   28, 24, 24, 28, 24, 24, 28, 24, 24, 28, 24, 24};
	
	public void ustawProgWektoryzacji(double wartosc) {
		wektoryzacjaKrzywej = new WektoryzacjaKrzywej(wartosc);
	}
	
	private List<int[]> wyszukujUzupelnienia(IImageChannel podzielony, int x, int y) {
		int minimalnyKoszt = 50;
		List<int[]> wynik = null;
		for (int i = 0; i < s1x.length; i++) {
			int x1 = x + s1x[i], y1 = y + s1y[i];
			int x2 = x + s2x[i], y2 = y + s2y[i];
			if (podzielony.inRange(x2, y2)) {
				float v1 = podzielony.getValue(x1, y1);
				float v2 = podzielony.getValue(x2, y2);
				if ((v1 > 0)&&(v1 < 1)&&(v2 == 2)) {
					if (cst[i] < minimalnyKoszt) {
						minimalnyKoszt = cst[i];
						wynik = new ArrayList<int[]>();
						wynik.add(new int[]{x1, y1});
						wynik.add(new int[]{x2, y2});
					}
				}
			}
		}
		return wynik;
	}
	
	private void uzupelnienieROI(IROI krzywa, IImageChannel podzielony) {
		List<int[]> punkty = krzywa.getROICoordinates();
		List<int[]> doDodania = new ArrayList<int[]>();
		for (int[] punkt : punkty) {
			List<int[]> wynik = this.wyszukujUzupelnienia(podzielony, punkt[0], punkt[1]);
			if (wynik != null) {
				doDodania.addAll(wynik);
				for (int[] nowy : wynik) {
					podzielony.setValue(nowy[0], nowy[1], 0.5f);
				}
			}
		}
		for (int[] punkt : doDodania) {
			krzywa.add(punkt[0], punkt[1]);
		}
	}
	
	private List<int[]> zjadaj(int x, int y, IROI krzywa, IImageChannel kanal, List<int[]> bs, int[] init) {
		if (!kanal.inRange(x, y)) {
			return null;
		}
		float v = kanal.getValue(x, y);
		if (v < 0.5f) {
			return null;
		}
		List<int[]> z = new ArrayList<int[]>();
		int[] p = new int[]{x, y};
		z.add(p);
		kanal.setValue(x, y, 0);
		if (init != null) {
			if (this.odleglosc.distI(p, init) > 1.5f) {
				kanal.setValue(init, 1);
				init = null;
			}
		}
		for (int[] b : bs) {
			if (b == init) {
				continue;
			}
			if (this.odleglosc.distI(b, p) < 1.5f) {
				z.add(b);
				return z;
			}
		}
		List<int[]> r = null;
		r = zjadaj(x + 1, y, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x - 1, y, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x, y + 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x, y - 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x + 1, y - 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x + 1, y + 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x - 1, y - 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		r = zjadaj(x - 1, y + 1, krzywa, kanal, bs, init); if (r != null) { z.addAll(r); return z; }
		return z;
	}
	
	public List<List<int[]>> wektoryzuj2(IImageChannel obraz) {
		//obraz = this.krawedzie(obraz);
		obraz = new BasicThreshold(obraz, 0.3f).apply();
		obraz = new FastParallelThinning(obraz, 0.5f).apply();
		
		return this.utworzWektory(obraz);
	}
	
	public List<List<int[]>> utworzWektory(IImageChannel obraz) {
		List<List<int[]>> wektory = new ArrayList<List<int[]>>();
		
		IROI[] krzywe = new SplitSeparate(obraz, 0.5f).apply();
		IImageChannel kanalKrzywej = new ImageChannel(obraz.getSize());
		for (IROI krzywa : krzywe) {
			this.rysowanieROI.apply(kanalKrzywej, krzywa);
			
			IROI terminalROI = new TerminationDetector().apply(krzywa, kanalKrzywej);
			IROI branchesROI = new BranchDetector().apply(krzywa, kanalKrzywej);
			
			List<int[]> terminal = terminalROI.getROICoordinates();
			List<int[]> branches = branchesROI.getROICoordinates();
			
			if ((terminal.size() == 0)&&(branches.size() == 0)) {
				wektory.add(krzywa.getROICoordinates());
			}
			for (int[] t : terminal) {
				/*for (int[] r2 : terminal) {
					kanalKrzywej.setValue(r2, 1);
				}*/
				for (int[] r2 : branches) {
					kanalKrzywej.setValue(r2, 1);
				}
				
				List<int[]> wynik = this.zjadaj(t[0], t[1], krzywa, kanalKrzywej, branches, null);
				if (wynik != null) {
					wektory.add(wynik);
				} else {
					System.out.println("Blad!");
				}
			}
			for (int[] b : branches) {
				/*for (int[] r2 : terminal) {
					kanalKrzywej.setValue(r2, 1);
				}*/
				for (int[] r2 : branches) {
					kanalKrzywej.setValue(r2, 1);
				}
				
				List<int[]> wynik = null;
				do {
					wynik = this.zjadaj(b[0], b[1], krzywa, kanalKrzywej, branches, b);
					if ((wynik != null)&&(wynik.size() > 1)) {
						wektory.add(wynik);
					}
				} while (wynik != null);
			}
			
			this.kasowanieROI.apply(kanalKrzywej, krzywa);
		}
		return wektory;
	}
	
	public List<List<int[]>> aproksymuj(List<List<int[]>> krzywe, IImageChannel kanal) {
		IImageChannel kanalFragmentu = new ImageChannel(kanal.getSize());
		List<List<int[]>> aproksymacje = new ArrayList<List<int[]>>();
		for (List<int[]> krzywa : krzywe) {
			IROI roi = new ROIPixels(kanal.getXSize(), kanal.getYSize(), krzywa);
			this.rysowanieROI.apply(kanalFragmentu, roi);
			List<int[]> aproksymacjaKrzywej = wektoryzacjaKrzywej.wektoryzuj(roi, kanalFragmentu);
			this.kasowanieROI.apply(kanalFragmentu, roi);
			aproksymacje.add(aproksymacjaKrzywej);
		}
		return aproksymacje;
	}
	
	private IImageChannel krawedzie(IImageChannel obraz) {
		IImageChannel wyjscie = new ImageChannel(obraz.getSize());
		for (int x = 1; x < obraz.getXSize() - 1; x++) {
			for (int y = 1; y < obraz.getYSize() - 1; y++) {
				float curVal = obraz.getValue(x, y);
				if (obraz.getValue(x - 1, y) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x + 1, y) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x, y - 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x, y + 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				
				if (obraz.getValue(x - 1, y - 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x + 1, y - 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x - 1, y + 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
				if (obraz.getValue(x + 1, y + 1) - curVal > 0.3f) { wyjscie.setValue(x, y, 1); }
			}
		}
		return wyjscie;
	}
	
	public List<List<int[]>> wektoryzuj(IImageChannel obraz) {
		//obraz = this.krawedzie(obraz);
		obraz = new BasicThreshold(obraz, 0.1f).apply();
		obraz = new FastParallelThinning(obraz, 0.5f).apply();
		IROI roi = new ROIFromThreshold(0.5f).apply(obraz);
		IROI terminalROI = new TerminationDetector().apply(roi, obraz);
		IROI branchesROI = new BranchDetector().apply(roi, obraz);
		
		List<int[]> terminal = terminalROI.getROICoordinates();
		List<int[]> branches = branchesROI.getROICoordinates();
				
		IImageChannel podzielony = new CloneChannel(obraz).apply();
		for (int[] point : branches) {
			List<int[]> maska = new ArrayList<int[]>();
			for (int x = -1; x <= 1; x++) {
				for (int y = -1; y <=1 ; y++) {
					maska.add(new int[]{point[0] + x, point[1] + y});
				}
			}
			this.maskowanie.apply(podzielony, maska);
		}
		
		List<List<int[]>> aproksymacjaObrazu = new ArrayList<List<int[]>>();
		IROI[] krzywe = new SplitSeparate(podzielony, 0.95f).apply();
		IImageChannel kanalFragmentu = new ImageChannel(obraz.getSize());
		for (IROI krzywa : krzywe) {
			for (int[] point : branches) {
				podzielony.setValue(point, 2);
			}
			this.uzupelnienieROI(krzywa, podzielony);
			this.rysowanieROI.apply(kanalFragmentu, krzywa);
			List<int[]> aproksymacjaKrzywej = wektoryzacjaKrzywej.wektoryzuj(krzywa, kanalFragmentu);
			aproksymacjaObrazu.add(aproksymacjaKrzywej);
			this.kasowanieROI.apply(kanalFragmentu, krzywa);
		}
		
		for (int i = 0; i < branches.size(); i++) {
			for (int j = i + 1; j < branches.size(); j++) {
				int[] bI = branches.get(i);
				int[] bJ = branches.get(j);
				float odl = odleglosc.distI(bI, bJ);
				if ((odl < 1.5)&&(odl > 0)) {
					List<int[]> krawedz = new ArrayList<int[]>();
					krawedz.add(bI);
					krawedz.add(bJ);
					aproksymacjaObrazu.add(krawedz);
				}
			}
		}
		return aproksymacjaObrazu;
	}
	
	public IRGBImage wyswietlAproksymacje(int[] wielkosc, List<List<int[]>> wektoryzacja) {
		IRGBImage obraz = new RGBImage(wielkosc);
		
		for (List<int[]> krzywa : wektoryzacja) {
			for (int i = 0; i < krzywa.size() - 1; i++) {
				int[] poczatek = krzywa.get(i);
				int[] koniec = krzywa.get(i + 1);
				new DrawLine(obraz.getGreenChannel(), poczatek, koniec, 1).apply();
				new DrawLine(obraz.getBlueChannel(), poczatek, koniec, 1).apply();
				new DrawLine(obraz.getRedChannel(), poczatek, koniec, 1).apply();
			}
			for (int i = 0; i < krzywa.size(); i++) {
				int[] punkt = krzywa.get(i);
				new DrawX(obraz.getBlueChannel(), punkt[0], punkt[1], 1, 1).apply();
			}
		}
		
		return obraz;
	}
	
	public IRGBImage wyswietlWektoryzacje(int[] wielkosc, List<List<int[]>> wektoryzacja) {
		IRGBImage obraz = new RGBImage(wielkosc);
		Random rand = new Random();
		
		for (List<int[]> krzywa : wektoryzacja) {
			float r = rand.nextFloat();
			float g = rand.nextFloat();
			float b = rand.nextFloat();
			for (int i = 0; i < krzywa.size(); i++) {
				int[] punkt = krzywa.get(i);
				obraz.getRedChannel().setValue(punkt, r);
				obraz.getGreenChannel().setValue(punkt, g);
				obraz.getBlueChannel().setValue(punkt, b);
			}
		}
		
		return obraz;
	}
	
	public IImageChannel przygotujKanal(IRGBImage image) {
		IGSImage grayImage = new RGBtoGSConverter(image).apply();
		IImageChannel kanal = grayImage.getGrayChannel();
		double srednia = new ImageChannelMean().applyDouble(kanal);
		if (srednia > 0.5) {
			kanal = new Invert(kanal).apply();
		}
		return kanal;
	}
	
	public boolean zapiszAproksymacje(List<List<int[]>> aproksymacja, File plik) {
		return this.zapisGrafu.zapisz(aproksymacja, plik);
	}
	
	public void przetworz(String[] args) {
		System.out.println("Image Vectorization, ver. 1.0, release date: 2010.10.04.");
		if (args.length < 1) {
			System.out.println("Author: Mariusz Paradowski, Wroclaw University of Technology");
			System.out.println("");
			System.out.println("ImageVectorization.jar input-image.png [threshold] [-draw]");
			System.out.println();
			System.out.println("   input-image - the image to be processed, should be");
			System.out.println("                 an edge or skeleton image");
			System.out.println("                 suggested input format is 24 bit PNG");
			System.out.println("                 although other formats may work as well");
			System.out.println("   threshold   - defines the approximation quality");
			System.out.println("                 the lower value, the better quality");
			System.out.println("                 default value: 0.5");
			System.out.println("   -draw       - drawing switch, draws the output image");
			System.out.println();
			System.out.println("The output file is a text file containing an edge graph.");
			System.out.println("First line of the output contains the number n of graph nodes.");
			System.out.println("Subsequent n lines contain image coordinates of each node.");
			System.out.println("Last n lines contain the definition of graph edges (for each node).");
			System.out.println("All nodes have 0->based index.");
			System.out.println("All edges are duplicated, both i --> j and j --> i edges are present.");
			return;
		} else {
			System.out.println("Call without parameters to see manual.");
		}
		File plik = new File(args[0]);
		if (!plik.exists()) {
			System.out.println("Input imageobraz = this.krawedzie(obraz); does not exists: " + plik.getName());
			return;
		}
		IRGBImage image = new GenericLoader(plik).apply();
		if (image == null) {
			System.out.println("Error on loading input image. Image format should be 24 bit PNG.");
			return;
		}
		double progWektoryzacji = 0.5;
		if (args.length > 1) {
			try {
				progWektoryzacji = Double.parseDouble(args[1]);
				progWektoryzacji = Math.max(0, progWektoryzacji);
			} catch (Exception ex) {
				System.out.println("Error on threshold parsing. Assumes threshold " + progWektoryzacji + ".");
			};
		}
		boolean rysuj = false;
		if (args.length > 2) {
			if ("-draw".equals(args[2])) {
				rysuj = true;
			}
		}
		System.out.println("Processing file: " + plik.getName());
		System.out.println("Threshold      : " + progWektoryzacji);
		System.out.println("Output drawing : " + rysuj);
		this.ustawProgWektoryzacji(progWektoryzacji);
		System.out.println("Preparing input image...");
		IImageChannel kanal = this.przygotujKanal(image);
		System.out.println("Calculating image vectorization...");
		List<List<int[]>> wektoryzacja = this.wektoryzuj2(kanal);
		System.out.println("Calculating curve approximation...");
		List<List<int[]>> aproksymacja = this.aproksymuj(wektoryzacja, kanal);
		File wyjscieWektory = new File(plik.getPath() + ".vector");
		System.out.println("Storing vector approximation...");
		this.zapiszAproksymacje(aproksymacja, wyjscieWektory);
		if (rysuj) {
			System.out.println("Drawing vector approximation...");
			IRGBImage obrazWektoryzacja = this.wyswietlWektoryzacje(image.getSize(), wektoryzacja);
			File wyjscieWektoryzacja = new File(plik.getPath() + ".vector.png");
			new GenericSaver(obrazWektoryzacja, wyjscieWektoryzacja).apply();
			IRGBImage obrazAproksymacja = this.wyswietlAproksymacje(image.getSize(), aproksymacja);
			File wyjscieAproksymacja = new File(plik.getPath() + ".approx.png");
			new GenericSaver(obrazAproksymacja, wyjscieAproksymacja).apply();
		}
	}
	
	public static void main(String[] args) {
		WektoryzacjaObrazu w = new WektoryzacjaObrazu();
		File folder = new File("/home/mariusz/dane/nekst/diagramy/test");
		List<File> obrazy = new ImageFolderScanner().apply(folder);
		for (File obraz : obrazy) {
			w.przetworz(new String[]{obraz.getPath(), "0.5", "-draw"});
		}
	}
}