<?php

namespace Rubix\ML\NeuralNet\Layers;

use Tensor\Matrix;
use Rubix\ML\Deferred;
use Rubix\ML\NeuralNet\Optimizers\Optimizer;
use Rubix\ML\NeuralNet\Initializers\Constant;
use Rubix\ML\NeuralNet\Parameter;
use Rubix\ML\NeuralNet\Initializers\Initializer;
use Rubix\ML\Exceptions\RuntimeException;
use Generator;

/**
 * PReLU
 *
 * Parametric Rectified Linear Units are leaky rectifiers whose leakage coefficients
 * are learned during training.
 *
 * References:
 * [1] K. He et al. (2015). Delving Deep into Rectifiers: Surpassing Human-Level
 * Performance on ImageNet Classification.
 *
 * @category    Machine Learning
 * @package     Rubix/ML
 * @author      Andrew DalPino
 */
class PReLU implements Hidden, Parametric
{
    /**
     * The initializer of the alpha (leakage) parameter.
     *
     * @var Initializer
     */
    protected Initializer $initializer;

    /**
     * The width of the layer.
     *
     * @var positive-int|null
     */
    protected ?int $width = null;

    /**
     * The parameterized leakage coefficients.
     *
     * @var Parameter|null
     */
    protected ?Parameter $alpha = null;

    /**
     * The memoized input matrix.
     *
     * @var Matrix|null
     */
    protected ?Matrix $input = null;

    /**
     * @param Initializer|null $initializer
     */
    public function __construct(?Initializer $initializer = null)
    {
        $this->initializer = $initializer ?? new Constant(0.25);
    }

    /**
     * Return the width of the layer.
     *
     * @internal
     *
     * @throws RuntimeException
     * @return positive-int
     */
    public function width() : int
    {
        if ($this->width === null) {
            throw new RuntimeException('Layer has not been initialized.');
        }

        return $this->width;
    }

    /**
     * Initialize the layer with the fan in from the previous layer and return
     * the fan out for this layer.
     *
     * @internal
     *
     * @param positive-int $fanIn
     * @return positive-int
     */
    public function initialize(int $fanIn) : int
    {
        $fanOut = $fanIn;

        $alpha = $this->initializer->initialize(1, $fanOut)->columnAsVector(0);

        $this->width = $fanOut;
        $this->alpha = new Parameter($alpha);

        return $fanOut;
    }

    /**
     * Compute a forward pass through the layer.
     *
     * @internal
     *
     * @param Matrix $input
     * @return Matrix
     */
    public function forward(Matrix $input) : Matrix
    {
        $this->input = $input;

        return $this->activate($input);
    }

    /**
     * Compute an inferential pass through the layer.
     *
     * @internal
     *
     * @param Matrix $input
     * @return Matrix
     */
    public function infer(Matrix $input) : Matrix
    {
        return $this->activate($input);
    }

    /**
     * Calculate the gradient and update the parameters of the layer.
     *
     * @internal
     *
     * @param Deferred $prevGradient
     * @param Optimizer $optimizer
     * @throws RuntimeException
     * @return Deferred
     */
    public function back(Deferred $prevGradient, Optimizer $optimizer) : Deferred
    {
        if (!$this->alpha) {
            throw new RuntimeException('Layer has not been initialized.');
        }

        if (!$this->input) {
            throw new RuntimeException('Must perform forward pass'
                . ' before backpropagating.');
        }

        $dOut = $prevGradient();

        $dIn = $this->input->clipUpper(0.0);

        $dAlpha = $dOut->multiply($dIn)->sum();

        $this->alpha->update($dAlpha, $optimizer);

        $input = $this->input;

        $this->input = null;

        return new Deferred([$this, 'gradient'], [$input, $dOut]);
    }

    /**
     * Calculate the gradient for the previous layer.
     *
     * @internal
     *
     * @param Matrix $input
     * @param Matrix $dOut
     * @return Matrix
     */
    public function gradient($input, $dOut) : Matrix
    {
        return $this->differentiate($input)->multiply($dOut);
    }

    /**
     * Return the parameters of the layer.
     *
     * @internal
     *
     * @throws \RuntimeException
     * @return Generator<Parameter>
     */
    public function parameters() : Generator
    {
        if (!$this->alpha) {
            throw new RuntimeException('Layer has not been initialized.');
        }

        yield 'alpha' => $this->alpha;
    }

    /**
     * Restore the parameters in the layer from an associative array.
     *
     * @internal
     *
     * @param Parameter[] $parameters
     */
    public function restore(array $parameters) : void
    {
        $this->alpha = $parameters['alpha'];
    }

    /**
     * Compute the leaky ReLU activation function and return a matrix.
     *
     * @param Matrix $input
     * @throws RuntimeException
     * @return Matrix
     */
    protected function activate(Matrix $input) : Matrix
    {
        if (!$this->alpha) {
            throw new RuntimeException('Layer has not been initialized.');
        }

        $alphas = $this->alpha->param()->asArray();

        $computed = [];

        foreach ($input as $i => $row) {
            $alpha = $alphas[$i];

            $activations = [];

            foreach ($row as $value) {
                $activations[] = $value > 0.0
                    ? $value
                    : $alpha * $value;
            }

            $computed[] = $activations;
        }

        return Matrix::quick($computed);
    }

    /**
     * Calculate the derivative of the activation function at a given output.
     *
     * @param Matrix $input
     * @throws RuntimeException
     * @return Matrix
     */
    protected function differentiate(Matrix $input) : Matrix
    {
        if (!$this->alpha) {
            throw new RuntimeException('Layer has not been initialized.');
        }

        $alphas = $this->alpha->param()->asArray();

        $gradient = [];

        foreach ($input as $i => $row) {
            $alpha = $alphas[$i];

            $derivative = [];

            foreach ($row as $value) {
                $derivative[] = $value > 0.0 ? 1.0 : $alpha;
            }

            $gradient[] = $derivative;
        }

        return Matrix::quick($gradient);
    }

    /**
     * Return the string representation of the object.
     *
     * @internal
     *
     * @return string
     */
    public function __toString() : string
    {
        return "PReLU (initializer: {$this->initializer})";
    }
}