<?php

declare(strict_types=1);

/**
 * SPDX-FileCopyrightText: 2018 Nextcloud GmbH and Nextcloud contributors
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

namespace OCA\SuspiciousLogin\Command;

use OCA\SuspiciousLogin\Exception\InsufficientDataException;
use OCA\SuspiciousLogin\Exception\ServiceException;
use OCA\SuspiciousLogin\Service\DataLoader;
use OCA\SuspiciousLogin\Service\Ipv4Strategy;
use OCA\SuspiciousLogin\Service\IpV6Strategy;
use OCA\SuspiciousLogin\Service\MLP\Trainer;
use OCA\SuspiciousLogin\Service\ModelStore;
use OCA\SuspiciousLogin\Service\TrainingDataConfig;
use Symfony\Component\Console\Command\Command;
use Symfony\Component\Console\Input\InputInterface;
use Symfony\Component\Console\Input\InputOption;
use Symfony\Component\Console\Output\OutputInterface;
use function extension_loaded;
use function time;

class Train extends Command {
	use ModelStatistics;

	/** @var DataLoader */
	private $loader;

	/** @var Trainer */
	private $trainer;

	/** @var ModelStore */
	private $store;

	public function __construct(DataLoader $loader,
		Trainer $optimizer,
		ModelStore $store) {
		parent::__construct('suspiciouslogin:train');
		$this->trainer = $optimizer;
		$this->loader = $loader;
		$this->store = $store;

		$this->addOption(
			'epochs',
			'e',
			InputOption::VALUE_OPTIONAL,
			'number of epochs to train'
		);
		$this->addOption(
			'layers',
			'l',
			InputOption::VALUE_OPTIONAL,
			'number of hidden layers'
		);
		$this->addOption(
			'shuffled',
			null,
			InputOption::VALUE_OPTIONAL,
			'ratio of shuffled negative samples'
		);
		$this->addOption(
			'random',
			null,
			InputOption::VALUE_OPTIONAL,
			'ratio of random negative samples'
		);
		$this->addOption(
			'learn-rate',
			null,
			InputOption::VALUE_OPTIONAL,
			'learning rate'
		);
		$this->addOption(
			'validation-threshold',
			null,
			InputOption::VALUE_OPTIONAL,
			'determines how much of the most recent data is used for validation. the default is one week'
		);
		$this->addOption(
			'max-age',
			null,
			InputOption::VALUE_OPTIONAL,
			'determines the maximum age of test data'
		);
		$this->addOption(
			'now',
			null,
			InputOption::VALUE_OPTIONAL,
			'overwrite the current time',
			time()
		);
		$this->addOption(
			'v6',
			null,
			InputOption::VALUE_NONE,
			'train with IPv6 data'
		);
		$this->addOption(
			'dry-run',
			null,
			InputOption::VALUE_NONE,
			"train but don't persist the model"
		);
		$this->addOption(
			'now',
			null,
			InputOption::VALUE_OPTIONAL,
			'the current time as timestamp',
			time()
		);
		$this->registerStatsOption();
	}

	protected function execute(InputInterface $input, OutputInterface $output): int {
		$strategy = $input->getOption('v6') ? new IpV6Strategy() : new Ipv4Strategy();
		$config = $strategy->getDefaultMlpConfig();
		if ($input->getOption('epochs') !== null) {
			$config = $config->setEpochs((int)$input->getOption('epochs'));
		}
		if ($input->getOption('layers') !== null) {
			$config = $config->setLayers((int)$input->getOption('layers'));
		}
		if ($input->getOption('shuffled') !== null) {
			$config = $config->setShuffledNegativeRate((float)$input->getOption('shuffled'));
		}
		if ($input->getOption('random') !== null) {
			$config = $config->setRandomNegativeRate((float)$input->getOption('random'));
		}
		if ($input->getOption('learn-rate') !== null) {
			$config = $config->setLearningRate((float)$input->getOption('learn-rate'));
		}

		$trainingDataConfig = TrainingDataConfig::default((int)$input->getOption('now'));
		if ($input->getOption('validation-threshold') !== null) {
			$trainingDataConfig = $trainingDataConfig->setThreshold((int)$input->getOption('validation-threshold'));
		}
		if ($input->getOption('max-age') !== null) {
			$trainingDataConfig = $trainingDataConfig->setMaxAge((int)$input->getOption('max-age'));
		}
		if ($input->getOption('now') !== null) {
			$trainingDataConfig = $trainingDataConfig->setNow((int)$input->getOption('now'));
		}

		try {
			if (extension_loaded('xdebug')) {
				$output->writeln('<comment>XDebug is active. This will slow down the training process.</comment>');
			}
			$output->writeln('Using ' . $strategy::getTypeName() . ' strategy');

			$collectedData = $this->loader->loadTrainingAndValidationData(
				$trainingDataConfig,
				$strategy
			);
			$data = $this->loader->generateRandomShuffledData(
				$collectedData,
				$config,
				$strategy
			);
			$result = $this->trainer->train(
				$config,
				$data,
				$strategy
			);
			$this->printModelStatistics($result->getModel(), $input, $output);
			if (!$input->getOption('dry-run')) {
				$this->store->persist(
					$result->getClassifier(),
					$result->getModel()
				);
				$output->writeln('<info>Model and estimator persisted.</info>');
			}
		} catch (InsufficientDataException $ex) {
			$output->writeln('<info>Not enough data, try again later (<error>' . $ex->getMessage() . '</error>)</info>');
			return 1;
		} catch (ServiceException $ex) {
			$output->writeln('<error>Could not train a model: ' . $ex->getMessage() . '</error>');
			return 1;
		}
		return 0;
	}
}
