/*
	Source Code taken from here: 
	https://github.com/noahvanhoucke/GalaxyZooChallenge/blob/master/CentralPixelBenchmark.go

	Edited and Modified by:
		Rohit Kumar Jha
		Irfan Hudda
	*/

package main

import (
	"bufio"
	"encoding/csv"
	"fmt"
	"image/jpeg"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func GetImageRGB(filename string) (int, int, int, int) {
	/* Given an image file:
	   - open file
	   - finds central pixel
	   - returns RGB values of central pixel
	*/
	// open image file
	file, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		println("Error opening file: ", err)
		os.Exit(-1)
	}
	defer file.Close()

	// reader takes the data from the file and puts into a buffer; used for optimization
	reader := bufio.NewReader(file)
	image, err := jpeg.Decode(reader)
	if err != nil {
		println("Error opening image file: ", err)
		os.Exit(-1)
	}

	// find central pixel, query 10x10 patch around it, average RGB results over patch
	center_x := (image.Bounds().Max.X - image.Bounds().Min.X)/2
	center_y := (image.Bounds().Max.Y - image.Bounds().Min.Y)/2
	avR := 0
	avG := 0
	avB := 0
	avA := 0
	count := 0
	for i := center_x-5; i < center_x+5; i++ {
		for j := center_y-5; j < center_y+5; j++ {
			r,g,b,a := image.At(i,j).RGBA()
			avR += int(r)
			avG += int(g)
			avB += int(b)
			avA += int(a)
			count += 1
		}
	}
	return avR/count, avG/count, avB/count, avA/count
}

func AssignClassValues(splitrow []string) []float64 {
	/* Converts the Class probability values read from file from 
	   string to float. */

	floatrow := make([]float64, 37)
	for i := 0; i < len(splitrow)-1; i++ {
		val, err := strconv.ParseFloat(splitrow[i], 64)
		floatrow[i] = val
		if err != nil {
			println("Error converting string to float ", err)
			os.Exit(-1)
		}
	}
	return floatrow
}

func GetTrainingSolutions(filename string) (map[string][]float64, []string) {
	/* Reads the training solutions file. Returns dictionary with GalaxyId
	   as the key and values as an array of float presenting the Class probability values. */

	trainingSolutions := make(map[string][]float64)
	trainFile, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		println("Error opening file: ", err)
		os.Exit(-1)
	}
	defer trainFile.Close()
	scanner := bufio.NewScanner(trainFile)
	scanner.Scan()
	header := scanner.Text()
	for scanner.Scan() {
		row := scanner.Text()
		splitrow := strings.Split(row, ",")
		values := AssignClassValues(splitrow[1:len(splitrow)])
		trainingSolutions[splitrow[0]] = values
	}
	headerRow := strings.Split(header, ",")
	return trainingSolutions, headerRow
}

func GetGalaxyRGB(filepathname string) map[string][3]int {
	/* Use Glob to get a list of all files in the images directory.
	   Return map with key = galaxyId, value = RGB values of central pixel. */

	listFiles, err := filepath.Glob(filepathname)
	if err != nil {
		println("Error using Glob")
		os.Exit(-1)
	}

	// Get the RGB values for the central pixel for each image
	println("GetGalaxyRGB(): Processing file ", filepathname)
	galaxyRGB := make(map[string][3]int)
	for i := 0; i < len(listFiles); i++ {
		galaxyId := strings.Split(listFiles[i], "\\")[2]
		galaxyId = strings.Split(galaxyId, ".")[0]
		r, g, b, _ := GetImageRGB(listFiles[i])
		// Print to stdout while processing files
		if i%5000 == 0 {
			fmt.Printf("i = %v, Galaxy Id = %v\n", i, galaxyId)
		}
		galaxyRGB[galaxyId] = [3]int{r, g, b}
	}

	if len(galaxyRGB) != len(listFiles) {
		println("GetGalaxyRGB(): Missing galaxy somewhere. galaxyRGB and listFiles have different lengths: ", len(galaxyRGB), len(listFiles))
		os.Exit(-1)
	}
	return galaxyRGB
}

func AverageGalaxySolutions(galaxyList []string, trainingSolutions map[string][]float64) []float64 {
	/* Average the training solutions for all galaxies in the galaxyList */

	var avSolutions []float64
	NProbabilities := len(trainingSolutions[galaxyList[0]])
	if NProbabilities != 37 {
		println("AverageGalaxySolutions(): NProbabilities is not 37, it is ", NProbabilities)
		os.Exit(-1)
	}

	// If only 1 galaxy, no average taken
	if len(galaxyList) == 1 {
		avSolutions = trainingSolutions[galaxyList[0]]
	} else {
		// Else average over all galaxies in the galaxyList
		avSolutions = trainingSolutions[galaxyList[0]]
		for i := 1; i < len(galaxyList); i++ {
			for j := 0; j < NProbabilities; j++ {
				avSolutions[j] += trainingSolutions[galaxyList[i]][j]
			}
		}
		for j := 0; j < NProbabilities; j++ {
			avSolutions[j] /= float64(len(galaxyList))
		}
	}
	return avSolutions
}

func GetGalaxyClusters(trainingGalaxyRGB map[string][3]int, hashFactor int) map[int][]string {
	/* Cluster training galaxies by central pixdel RGB values.
	Divide RGB values by the average intensity to work with just the color.
	Use a hashFactor to do a quick and dirty clustering. */

	galaxyClusters := make(map[int][]string)
	for galaxyId, RGB := range trainingGalaxyRGB {
		avIntensity := int((RGB[0] + RGB[1] + RGB[2]) / 3)
		RGB[0] /= avIntensity
		RGB[1] /= avIntensity
		RGB[2] /= avIntensity
		key := RGB[0]*hashFactor*hashFactor + RGB[1]*hashFactor + RGB[2]
		if _, ok := galaxyClusters[key]; ok {
			galaxyClusters[key] = append(galaxyClusters[key], galaxyId)
		} else {
			galaxyClusters[key] = []string{galaxyId}
		}
	}
	if len(galaxyClusters) < 1 {
		println("GetGalaxyClusters(): Created zero galaxy clusters.")
		os.Exit(-1)
	}
	return galaxyClusters
}

func GetSolutionsForGalaxyClusters(galaxyClusters map[int][]string, trainingSolutions map[string][]float64) map[int][]float64 {
	/* Averages the training solution values over each galaxy cluster.
	   Arguments:
	     - trainingSolutions: For each galaxyId, the solutions -- the 37 probability values.
	     - galaxyClusters: Central pixel has been clustered and mapped to a unqiue value.
	          map's key is the unique value; map's value is a list of galaxyId's for that cluster
	   Returns:
	     - galaxyClusterSolutions: probability values for each galaxy cluster. Is average
	     over a cluster's solutions
	*/

	galaxyClusterSolutions := make(map[int][]float64)
	for clusterId, galaxyList := range galaxyClusters {
		galaxyClusterSolutions[clusterId] = AverageGalaxySolutions(galaxyList, trainingSolutions)
	}
	return galaxyClusterSolutions
}

func BuildPredictionRow(galaxyId string, pred []float64) []string {
	/* Build the string to be written to file using the csv library. */

	var line []string
	line = append(line, galaxyId)
	for i := 0; i < len(pred); i++ {
		line = append(line, strconv.FormatFloat(pred[i], 'f', -1, 64))
	}

	if len(line) != 38 {
		println("BuildPredictionRow(): line does not have 38 elements, it has: ", len(line))
		os.Exit(-1)
	}
	return line
}

func CreateCentralPixelBenchmark(galaxyClusterSolutions map[int][]float64,
	testGalaxyRGB map[string][3]int, testPredictionFile string, hashFactor int, headerRow []string) {
	/* Central Pixel Benchmark: using the average values for solutions for galaxies clustered by the
	   central pixel. If no match, then use 0.0 for all values. */

	testPredictions := make(map[string][]float64)
	for galaxyId, RGB := range testGalaxyRGB {
		avIntensity := int((RGB[0] + RGB[1] + RGB[2]) / 3)
		rgbCluster := hashFactor*hashFactor*(RGB[0]/avIntensity) + hashFactor*(RGB[1]/avIntensity) + RGB[2]/avIntensity
		if _, ok := galaxyClusterSolutions[rgbCluster]; ok {
			testPredictions[galaxyId] = galaxyClusterSolutions[rgbCluster]
		} else {
			println("CreateCentralPixelBenchmark(): Key not found for this GalaxyId, rgbCluster = ", galaxyId, rgbCluster)
			testPredictions[galaxyId] = []float64{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}
		}
	}

	// Write test predictions to file
	fo, err := os.Create(testPredictionFile)
	if err != nil {
		println("CreateCentralPixelBenchmark(): Error opening file: ", err)
		os.Exit(-1)
	}
	defer fo.Close()
	wr := csv.NewWriter(fo)

	// write header to file
	err = wr.Write(headerRow)
	println("Lengths of testPredictions and testGalaxyRGB: ", len(testPredictions), len(testGalaxyRGB))
	count := 0
	if len(testPredictions) == len(testGalaxyRGB) {
		for galaxyId, _ := range testGalaxyRGB {
			if len(testPredictions[galaxyId]) != 37 {
				println("Do not have 37 predictions for galaxy ", galaxyId)
				os.Exit(-1)
			}
			err = wr.Write(BuildPredictionRow(galaxyId, testPredictions[galaxyId]))
			count += 1
		}
		wr.Flush()
	} else {
		println("Lengths don't match -- missing galaxy somewhere")
		os.Exit(-1)
	}
	println("Number of predictions made = ", count)
}

func main() {

	// Get RGB values for central pixel in training and test images
	trainingImagesPath := "images_training/*.jpg"
	testImagesPath := "images_test/*.jpg"
	trainingGalaxyRGB := GetGalaxyRGB(trainingImagesPath)
	testGalaxyRGB := GetGalaxyRGB(testImagesPath)

	trainingSolutions, headerRow := GetTrainingSolutions("solutions_training.csv")

	hashFactor := 10
	galaxyClusters := GetGalaxyClusters(trainingGalaxyRGB, hashFactor)
	galaxyClusterSolutions := GetSolutionsForGalaxyClusters(galaxyClusters, trainingSolutions)

	// Given test image central pixel RGB value, find matching galaxy cluster and
	// use the average galaxy cluster solution for the test image's solution
	testPredictionFile := "lastrun.csv"
	CreateCentralPixelBenchmark(galaxyClusterSolutions, testGalaxyRGB, testPredictionFile, hashFactor, headerRow)

	println("\nEnd Program. Success!")
}
