package main

import "core:container/bit_array"
import "core:encoding/csv"
import "core:fmt"
import "core:math"
import "core:math/rand"
import "core:os"
import "core:slice"
import "core:strconv"
import "core:strings"

OUTPUT_FILE :: "output/data.csv"
DATA_FILE :: "res/knapPI_12_500_1000_82.csv"
NUMBER_OF_ITEMS :: 500
CAPACITY :: 280785
POPULATION_SIZE :: 100
GENERATIONS :: 100
TOURNAMENT_SIZE :: 3
CROSSOVER_RATE :: 0.8
MUTATION_RATE :: 0.01

Item :: struct {
	profit, weight: int,
}

Chromosome :: ^bit_array.Bit_Array
Population :: [POPULATION_SIZE]Chromosome

items: [NUMBER_OF_ITEMS]Item

Data :: struct {
	best, worst: int,
	mean:        f32,
}

stats: [GENERATIONS]Data

read_data :: proc(file: string) -> (res: [NUMBER_OF_ITEMS]Item, ok := true) {
	data := string(os.read_entire_file(file) or_return)
	defer delete(data)

	for line, idx in strings.split_lines(strings.trim_space(data))[1:] {
		s := strings.split(line, ",")
		res[idx] = {strconv.parse_int(s[1]) or_return, strconv.parse_int(s[2]) or_return}
	}
	return
}

write_results :: proc(filename: string, stats: []Data) -> bool {
	handle, err := os.open(filename, os.O_CREATE | os.O_WRONLY | os.O_TRUNC, 0o644)
	if err != os.ERROR_NONE {return false}
	defer os.close(handle)

	w: csv.Writer
	csv.writer_init(&w, os.stream_from_handle(handle))

	csv.write(&w, []string{"generation", "best", "worst", "mean"})

	for stat, gen in stats {
		csv.write(
			&w,
			[]string {
				fmt.tprintf("%d", gen),
				fmt.tprintf("%d", stat.best),
				fmt.tprintf("%d", stat.worst),
				fmt.tprintf("%f", stat.mean),
			},
		)
	}

	csv.writer_flush(&w)
	return true
}


fitness :: proc(chrom: Chromosome) -> int {
	tot_profit, tot_weight := 0, 0
	for idx in 0 ..< bit_array.len(chrom) {
		if !bit_array.get(chrom, idx) {continue}
		tot_profit += items[idx].profit
		tot_weight += items[idx].weight
	}
	return tot_profit - 500 * max(tot_weight - CAPACITY, 0)
}

create_random_chromosome :: proc() -> Chromosome {
	chrom := bit_array.create(NUMBER_OF_ITEMS)
	for i in 0 ..< NUMBER_OF_ITEMS {
		bit_array.set(chrom, i, rand.int_max(2) == 1)
	}
	return chrom
}

copy_chromosome :: proc(src: Chromosome) -> Chromosome {
	dest := bit_array.create(NUMBER_OF_ITEMS)
	for i in 0 ..< NUMBER_OF_ITEMS {
		bit_array.set(dest, i, bit_array.get(src, i))
	}
	return dest
}

generate_population :: proc() -> Population {
	pop: Population
	for i in 0 ..< POPULATION_SIZE {
		pop[i] = create_random_chromosome()
	}
	return pop
}

destroy_population :: proc(pop: ^Population) {
	for chrom in pop {
		bit_array.destroy(chrom)
	}
}

evaluate_population :: proc(pop: ^Population) -> [POPULATION_SIZE]int {
	fitnesses: [POPULATION_SIZE]int
	for chrom, i in pop {
		fitnesses[i] = fitness(chrom)
	}
	return fitnesses
}

tournament_selection :: proc(
	pop: ^Population,
	fitnesses: []int,
	k := TOURNAMENT_SIZE,
) -> Chromosome {
	best_idx := rand.int_max(POPULATION_SIZE)
	best_fitness := fitnesses[best_idx]

	for _ in 1 ..< k {
		idx := rand.int_max(POPULATION_SIZE)
		if fitnesses[idx] > best_fitness {
			best_idx = idx
			best_fitness = fitnesses[idx]
		}
	}

	return pop[best_idx]
}

roulette_selection :: proc(pop: ^Population, fitnesses: []int) -> Chromosome {
	total_fitness := 0
	for f in fitnesses {
		total_fitness += max(f, 0)
	}

	if total_fitness == 0 {
		return pop[rand.int_max(POPULATION_SIZE)]
	}

	spin := rand.float32() * f32(total_fitness)
	running_sum := 0

	for fitness, idx in fitnesses {
		running_sum += max(fitness, 0)
		if f32(running_sum) >= spin {
			return pop[idx]
		}
	}

	return pop[POPULATION_SIZE - 1]
}

single_point_crossover :: proc(parent1, parent2: Chromosome) -> (child1, child2: Chromosome) {
	point := rand.int_max(NUMBER_OF_ITEMS)
	child1 = bit_array.create(NUMBER_OF_ITEMS)
	child2 = bit_array.create(NUMBER_OF_ITEMS)

	for i in 0 ..< NUMBER_OF_ITEMS {
		if i < point {
			bit_array.set(child1, i, bit_array.get(parent1, i))
			bit_array.set(child2, i, bit_array.get(parent2, i))
		} else {
			bit_array.set(child1, i, bit_array.get(parent2, i))
			bit_array.set(child2, i, bit_array.get(parent1, i))
		}
	}

	return
}

two_point_crossover :: proc(parent1, parent2: Chromosome) -> (child1, child2: Chromosome) {
	point1 := rand.int_max(NUMBER_OF_ITEMS)
	point2 := rand.int_max(NUMBER_OF_ITEMS)
	if point1 > point2 {
		point1, point2 = point2, point1
	}

	child1 = bit_array.create(NUMBER_OF_ITEMS)
	child2 = bit_array.create(NUMBER_OF_ITEMS)

	for i in 0 ..< NUMBER_OF_ITEMS {
		in_swap := i >= point1 && i < point2
		if in_swap {
			bit_array.set(child1, i, bit_array.get(parent2, i))
			bit_array.set(child2, i, bit_array.get(parent1, i))
		} else {
			bit_array.set(child1, i, bit_array.get(parent1, i))
			bit_array.set(child2, i, bit_array.get(parent2, i))
		}
	}

	return
}

uniform_crossover :: proc(
	parent1, parent2: Chromosome,
	prob := f32(0.5),
) -> (
	child1, child2: Chromosome,
) {
	child1 = bit_array.create(NUMBER_OF_ITEMS)
	child2 = bit_array.create(NUMBER_OF_ITEMS)

	for i in 0 ..< NUMBER_OF_ITEMS {
		if rand.float32() < prob {
			bit_array.set(child1, i, bit_array.get(parent1, i))
			bit_array.set(child2, i, bit_array.get(parent2, i))
		} else {
			bit_array.set(child1, i, bit_array.get(parent2, i))
			bit_array.set(child2, i, bit_array.get(parent1, i))
		}
	}

	return
}

bit_flip_mutation :: proc(chrom: Chromosome, mutation_rate := MUTATION_RATE) {
	for i in 0 ..< NUMBER_OF_ITEMS {
		if rand.float32() < f32(mutation_rate) {
			bit_array.set(chrom, i, !bit_array.get(chrom, i))
		}
	}
}

swap_mutation :: proc(chrom: Chromosome) {
	idx1 := rand.int_max(NUMBER_OF_ITEMS)
	idx2 := rand.int_max(NUMBER_OF_ITEMS)
	bit1 := bit_array.get(chrom, idx1)
	bit2 := bit_array.get(chrom, idx2)
	bit_array.set(chrom, idx1, bit2)
	bit_array.set(chrom, idx2, bit1)
}

inversion_mutation :: proc(chrom: Chromosome) {
	point1 := rand.int_max(NUMBER_OF_ITEMS)
	point2 := rand.int_max(NUMBER_OF_ITEMS)
	if point1 > point2 {
		point1, point2 = point2, point1
	}

	for i in 0 ..< (point2 - point1) / 2 {
		idx1 := point1 + i
		idx2 := point2 - i - 1
		bit1 := bit_array.get(chrom, idx1)
		bit2 := bit_array.get(chrom, idx2)
		bit_array.set(chrom, idx1, bit2)
		bit_array.set(chrom, idx2, bit1)
	}
}

create_offspring :: proc(population: ^Population, fitnesses: []int) -> Population {
	offspring: Population

	for i := 0; i < POPULATION_SIZE; i += 2 {
		parent1 := tournament_selection(population, fitnesses)
		parent2 := tournament_selection(population, fitnesses)

		child1, child2: Chromosome
		if rand.float32() < CROSSOVER_RATE {
			child1, child2 = single_point_crossover(parent1, parent2)
		} else {
			child1 = copy_chromosome(parent1)
			child2 = copy_chromosome(parent2)
		}

		bit_flip_mutation(child1)
		bit_flip_mutation(child2)

		offspring[i] = child1
		if i + 1 < POPULATION_SIZE {
			offspring[i + 1] = child2
		} else {
			bit_array.destroy(child2)
		}
	}

	return offspring
}

elitism_survivor_selection :: proc(
	parents: ^Population,
	offspring: ^Population,
	parent_fitnesses: []int,
	offspring_fitnesses: []int,
) -> Population {
	Index_Fitness :: struct {
		idx:       int,
		fitness:   int,
		is_parent: bool,
	}

	combined := make([]Index_Fitness, POPULATION_SIZE * 2)
	defer delete(combined)

	for i in 0 ..< POPULATION_SIZE {
		combined[i] = {i, parent_fitnesses[i], true}
		combined[POPULATION_SIZE + i] = {i, offspring_fitnesses[i], false}
	}

	slice.sort_by(combined[:], proc(a, b: Index_Fitness) -> bool {
		return a.fitness > b.fitness
	})

	survivors: Population
	for i in 0 ..< POPULATION_SIZE {
		entry := combined[i]
		source := parents[entry.idx] if entry.is_parent else offspring[entry.idx]
		survivors[i] = copy_chromosome(source)
	}

	return survivors
}

compute_stats :: proc(fitnesses: []int) -> Data {
	best, worst := math.min(int), math.max(int)
	sum := 0
	for f in fitnesses {
		best = max(best, f)
		worst = min(worst, f)
		sum += f
	}
	return {best, worst, f32(sum) / f32(len(fitnesses))}
}

run_ga :: proc() {
	population := generate_population()
	defer destroy_population(&population)

	best_fitness := math.min(int)
	best_generation := 0
	best_chromosome: Chromosome

	for gen in 0 ..< GENERATIONS {
		fitnesses := evaluate_population(&population)
		stats[gen] = compute_stats(fitnesses[:])

		for f, i in fitnesses {
			if f <= best_fitness {continue}
			best_fitness = f
			best_generation = gen
			bit_array.destroy(best_chromosome)
			best_chromosome = copy_chromosome(population[i])
			tot_profit, tot_weight := 0, 0
			for idx in 0 ..< bit_array.len(best_chromosome) {
				if !bit_array.get(best_chromosome, idx) {continue}
				tot_profit += items[idx].profit
				tot_weight += items[idx].weight
			}
			fmt.printfln(
				"Gen %d: Fitness=%d, Profit=%d, Weight=%d/%d",
				gen,
				f,
				tot_profit,
				tot_weight,
				CAPACITY,
			)
		}

		offspring := create_offspring(&population, fitnesses[:])
		offspring_fitnesses := evaluate_population(&offspring)

		new_population := elitism_survivor_selection(
			&population,
			&offspring,
			fitnesses[:],
			offspring_fitnesses[:],
		)

		destroy_population(&population)
		destroy_population(&offspring)

		population = new_population
	}

	fmt.printfln("\nFinal Best: Fitness=%d (Generation %d)", best_fitness, best_generation)
	fmt.println("this solution is the following bit-string:")
	for i in 0 ..< best_chromosome.length {
		b := bit_array.get(best_chromosome, i)
		fmt.print(i32(b))
	}
	fmt.println()
	write_results(OUTPUT_FILE, stats[:])
	fmt.println("successfully wrote data to", OUTPUT_FILE)
}

main :: proc() {
	data, ok := read_data(DATA_FILE)
	if !ok {
		fmt.eprintln("Failed to read data from", DATA_FILE)
		return
	}
	items = data

	fmt.println("Running Genetic Algorithm for Binary Knapsack Problem")
	fmt.printfln(
		"Items: %d, Capacity: %d, Population: %d, Generations: %d\n",
		NUMBER_OF_ITEMS,
		CAPACITY,
		POPULATION_SIZE,
		GENERATIONS,
	)

	run_ga()
}