"use strict";

import path from 'path';
import fs from 'fs';

import brain from 'brain.js';

class AITrainer {
	constructor({ ansi, settings, log, root_dir, GatewayRepo, DatasetFetcher }) {
		this.a = ansi;
		this.settings = settings;
		this.root_dir = root_dir;
		this.l = log;
		this.dataset_fetcher = DatasetFetcher;
		this.repo_gateway = GatewayRepo;
	}
	
	generate_neural_net() {
		let net = new brain.NeuralNetwork({
			hiddenLayers: this.settings.ai.network_arch,
			activation: "sigmoid"
		});
		
		
		return net;
	}
	
	async train_all() {
		let index = [];
		for(let gateway of this.repo_gateway.iterate()) {
			let filename = path.join(this.root_dir, "..", this.settings.ai.output_directory, `${gateway.id}`);
			console.log(filename);
			
			if(!fs.existsSync(path.dirname(filename)))
				await fs.promises.mkdir(path.dirname(filename), { recursive: true });
			
			if(!await this.train_gateway(gateway.id, filename)) {
				this.l.warn(`Warning: Failed to train AI for ${gateway.id}.`);
				continue;
			}
			
			index.push({
				id: gateway.id,
				latitude: gateway.latitude,
				longitude: gateway.longitude
			});
		}
		
		await fs.promises.writeFile(
			path.join(
				path.dirname(this.root_dir),
				this.settings.ai.output_directory,
				"index.json"
			),
			JSON.stringify({
				properties: {
					rssi_min: this.settings.ai.rssi_min,
					rssi_max:  this.settings.ai.rssi_max
				},
				index
			})
		);
	}
	
	/**
	 * Trains an AI to predict the coverage of a specific gateway.
	 * @param	{string}	gateway_id				The id of the gateway to train an AI for.
	 * @param	{string}	destination_filename	The absolute path to the file to serialise the trained to. Required because we can't serialise and return a TensorFlow model, it has to be sent somewhere because the API is backwards and upside-down :-/
	 * @return	{Promise}	A promise that resolves when training and serialisation is complete.
	 */
	async train_gateway(gateway_id, destination_filename) {
		this.l.log(`${this.a.fgreen}${this.a.hicol}Training AI for gateway ${gateway_id}${this.a.reset}`);
		let net = this.generate_neural_net();
		
		// let dataset_input = tf.data.generator(
		// 	this.dataset_fetcher.fetch_input.bind(this.dataset_fetcher, gateway_id)
		// );
		// let dataset_output = tf.data.generator(
		// 	this.dataset_fetcher.fetch_output.bind(this.dataset_fetcher, gateway_id)
		// );
		// 
		// let dataset = tf.data.zip({
		// 		xs: dataset_input,
		// 		ys: dataset_output
		// 	}).shuffle(this.settings.ai.batch_size * 4)
		// 	.batch(this.settings.ai.batch_size);
		// 
		let dataset = this.dataset_fetcher.fetch_all(gateway_id);
		
		let result = net.train(dataset, {
			iterations: this.settings.ai.epochs,
			errorThresh: this.settings.ai.error_threshold,
			
			learningRate: this.settings.ai.learning_rate,
			momentum: this.settings.ai.momentum,
			
			timeout: Infinity
		});
		
		await model.save(`file://${destination_filename}`);
		// console.log(result);
		
		return true;
	}
}

export default AITrainer;