/*
 * SPDX-License-Identifier: MIT
 */
package org.ta4j.core.indicators.supportresistance;

import static org.ta4j.core.indicators.IndicatorUtils.isInvalid;
import static org.ta4j.core.num.NaN.NaN;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

import org.ta4j.core.BarSeries;
import org.ta4j.core.Indicator;
import org.ta4j.core.indicators.CachedIndicator;
import org.ta4j.core.indicators.helpers.ClosePriceIndicator;
import org.ta4j.core.indicators.helpers.ConstantIndicator;
import org.ta4j.core.num.Num;
import org.ta4j.core.num.NumFactory;

/**
 * Base class for price cluster support and resistance indicators.
 *
 * <p>
 * Concrete implementations provide the tie-breaking behaviour for equally
 * popular clusters.
 *
 * @since 0.22.3
 */
public abstract class AbstractPriceClusterIndicator extends CachedIndicator<Num> {

    private final Indicator<Num> priceIndicator;
    private final Indicator<Num> weightIndicatorSource;
    private final int lookbackCount;
    private final Num tolerance;
    private transient Indicator<Num> weightIndicator;
    private transient Map<Integer, Integer> clusterIndexCache;
    private transient int lastPrunedCacheBeginIndex;

    /**
     * Constructor.
     *
     * @param priceIndicator the price source to cluster
     * @param lookbackCount  the number of bars to evaluate (non-positive for the
     *                       full history)
     * @param tolerance      the absolute tolerance for bucket membership
     * @since 0.22.3
     */
    protected AbstractPriceClusterIndicator(Indicator<Num> priceIndicator, int lookbackCount, Num tolerance) {
        this(priceIndicator, null, lookbackCount, tolerance);
    }

    /**
     * Constructor supporting a custom weight indicator.
     *
     * @param priceIndicator  the price source to cluster
     * @param weightIndicator optional weight indicator (defaults to unit weights
     *                        when {@code null})
     * @param lookbackCount   the number of bars to evaluate (non-positive for the
     *                        full history)
     * @param tolerance       the absolute tolerance for bucket membership
     * @since 0.22.3
     */
    protected AbstractPriceClusterIndicator(Indicator<Num> priceIndicator, Indicator<Num> weightIndicator,
            int lookbackCount, Num tolerance) {
        this(priceIndicator, weightIndicator, weightIndicator, lookbackCount, tolerance);
    }

    /**
     * Constructor supporting separate source and derived weight indicators.
     *
     * @param priceIndicator        the price source to cluster
     * @param weightIndicator       indicator used for weighting calculations
     * @param weightIndicatorSource logical weight indicator to expose for
     *                              serialization
     * @param lookbackCount         the number of bars to evaluate (non-positive for
     *                              the full history)
     * @param tolerance             the absolute tolerance for bucket membership
     * @since 0.22.3
     */
    protected AbstractPriceClusterIndicator(Indicator<Num> priceIndicator, Indicator<Num> weightIndicator,
            Indicator<Num> weightIndicatorSource, int lookbackCount, Num tolerance) {
        super(priceIndicator);
        this.priceIndicator = Objects.requireNonNull(priceIndicator, "priceIndicator must not be null");
        this.lookbackCount = lookbackCount;
        this.tolerance = Objects.requireNonNull(tolerance, "tolerance must not be null");
        BarSeries series = Objects.requireNonNull(priceIndicator.getBarSeries(),
                "indicator must reference a bar series");
        if (isInvalid(tolerance) || tolerance.isLessThan(series.numFactory().zero())) {
            throw new IllegalArgumentException("tolerance must be greater than or equal to zero");
        }
        Indicator<Num> resolvedWeight = weightIndicator;
        if (resolvedWeight == null) {
            resolvedWeight = new ConstantIndicator<>(series, series.numFactory().one());
        }
        Indicator<Num> resolvedSource = weightIndicatorSource != null ? weightIndicatorSource : resolvedWeight;
        if (resolvedWeight.getBarSeries() != series) {
            throw new IllegalArgumentException("weightIndicator must share the same bar series as priceIndicator");
        }
        if (resolvedSource.getBarSeries() != series) {
            throw new IllegalArgumentException(
                    "weightIndicatorSource must share the same bar series as priceIndicator");
        }
        this.weightIndicator = resolvedWeight;
        this.weightIndicatorSource = resolvedSource;
        this.clusterIndexCache = new ConcurrentHashMap<>();
        this.lastPrunedCacheBeginIndex = series.getBeginIndex() - 1;
    }

    /**
     * Constructor.
     *
     * @param series        the bar series backing the indicator
     * @param lookbackCount the number of bars to evaluate (non-positive for the
     *                      full history)
     * @param tolerance     the absolute tolerance for bucket membership
     * @since 0.22.3
     */
    protected AbstractPriceClusterIndicator(BarSeries series, int lookbackCount, Num tolerance) {
        this(new ClosePriceIndicator(series), lookbackCount, tolerance);
    }

    /**
     * Calculates the indicator value at the requested index.
     */
    @Override
    protected Num calculate(int index) {
        BarSeries series = getBarSeries();
        if (series == null || index < series.getBeginIndex()) {
            clusterIndexCache().put(index, -1);
            return NaN;
        }
        pruneClusterIndexCache(series);
        if (index < series.getBeginIndex() + getCountOfUnstableBars()) {
            clusterIndexCache().put(index, -1);
            return NaN;
        }

        int startIndex = computeStartIndex(index, series);
        List<PriceCluster> clusters = buildClusters(startIndex, index);
        PriceCluster bestCluster = selectBestCluster(clusters);

        if (bestCluster == null) {
            clusterIndexCache().put(index, -1);
            return NaN;
        }

        clusterIndexCache().put(index, bestCluster.getLastIndex());
        return bestCluster.getRepresentativePrice();
    }

    /**
     * Returns the unstable warmup covering source indicator warmup plus the
     * configured look-back window.
     *
     * @return number of unstable bars
     * @since 0.22.3
     */
    @Override
    public int getCountOfUnstableBars() {
        int componentUnstableBars = Math.max(priceIndicator.getCountOfUnstableBars(),
                weightIndicator().getCountOfUnstableBars());
        return componentUnstableBars + Math.max(0, lookbackCount - 1);
    }

    /**
     * Returns the index of the most recent price that belongs to the dominant
     * cluster at the requested bar.
     *
     * @param index the bar index
     * @return the most recent cluster member index or {@code -1} when no cluster is
     *         available
     * @since 0.22.3
     */
    public int getClusterIndex(int index) {
        getValue(index);
        return clusterIndexCache().getOrDefault(index, -1);
    }

    /**
     * Computes start index.
     */
    private int computeStartIndex(int index, BarSeries series) {
        if (lookbackCount <= 0) {
            return series.getBeginIndex();
        }
        int desiredStart = index - lookbackCount + 1;
        return Math.max(series.getBeginIndex(), desiredStart);
    }

    /**
     * Builds clusters.
     */
    private List<PriceCluster> buildClusters(int startIndex, int endIndex) {
        List<PriceCluster> clusters = new ArrayList<>();
        for (int i = startIndex; i <= endIndex; i++) {
            Num value = priceIndicator.getValue(i);
            if (isInvalid(value)) {
                continue;
            }
            Num weight = weightIndicator().getValue(i);
            if (isInvalid(weight) || !weight.isPositive()) {
                continue;
            }
            boolean assigned = false;
            for (PriceCluster cluster : clusters) {
                if (cluster.tryInclude(value, weight, i, tolerance)) {
                    assigned = true;
                    break;
                }
            }
            if (!assigned) {
                clusters.add(new PriceCluster(value, weight, i));
            }
        }
        return clusters;
    }

    /**
     * Selects best cluster.
     */
    private PriceCluster selectBestCluster(List<PriceCluster> clusters) {
        PriceCluster best = null;
        for (PriceCluster cluster : clusters) {
            if (best == null || cluster.getTotalWeight().isGreaterThan(best.getTotalWeight())) {
                best = cluster;
                continue;
            }
            if (cluster.getTotalWeight().isEqual(best.getTotalWeight())) {
                if (cluster.getCount() > best.getCount()) {
                    best = cluster;
                    continue;
                }
                if (cluster.getCount() == best.getCount()) {
                    Num clusterPrice = cluster.getRepresentativePrice();
                    Num bestPrice = best.getRepresentativePrice();
                    if (preferLowerPriceOnTie()) {
                        if (clusterPrice.isLessThan(bestPrice)) {
                            best = cluster;
                            continue;
                        }
                    } else {
                        if (clusterPrice.isGreaterThan(bestPrice)) {
                            best = cluster;
                            continue;
                        }
                    }
                    if (clusterPrice.isEqual(bestPrice) && cluster.getLastIndex() > best.getLastIndex()) {
                        best = cluster;
                    }
                }
            }
        }
        return best;
    }

    /**
     * Indicates whether tied clusters should favour the lower price.
     *
     * @return {@code true} to prefer lower prices, {@code false} to prefer higher
     *         prices
     * @since 0.22.3
     */
    protected abstract boolean preferLowerPriceOnTie();

    private static final class PriceCluster {
        private Num weightedSum;
        private Num totalWeight;
        private int count;
        private Num representativePrice;
        private int lastIndex;

        /**
         * Implements price cluster.
         */
        private PriceCluster(Num value, Num weight, int index) {
            this.weightedSum = value.multipliedBy(weight);
            this.totalWeight = weight;
            this.count = 1;
            this.representativePrice = value;
            this.lastIndex = index;
        }

        /**
         * Implements try include.
         */
        private boolean tryInclude(Num value, Num weight, int index, Num tolerance) {
            if (value.minus(representativePrice).abs().isLessThanOrEqual(tolerance)) {
                weightedSum = weightedSum.plus(value.multipliedBy(weight));
                totalWeight = totalWeight.plus(weight);
                count++;
                representativePrice = weightedSum.dividedBy(totalWeight);
                lastIndex = index;
                return true;
            }
            return false;
        }

        /**
         * Returns the total weight.
         */
        private Num getTotalWeight() {
            return totalWeight;
        }

        /**
         * Returns the count.
         */
        private int getCount() {
            return count;
        }

        /**
         * Returns the representative price.
         */
        private Num getRepresentativePrice() {
            return representativePrice;
        }

        /**
         * Returns the last index.
         */
        private int getLastIndex() {
            return lastIndex;
        }
    }

    /**
     * Implements weight indicator.
     */
    private Indicator<Num> weightIndicator() {
        if (weightIndicator == null) {
            weightIndicator = weightIndicatorSource;
        }
        return weightIndicator;
    }

    /**
     * Implements cluster index cache.
     */
    private Map<Integer, Integer> clusterIndexCache() {
        if (clusterIndexCache == null) {
            clusterIndexCache = new ConcurrentHashMap<>();
        }
        return clusterIndexCache;
    }

    /**
     * Prunes cluster index cache.
     */
    private void pruneClusterIndexCache(BarSeries series) {
        int beginIndex = series.getBeginIndex();
        if (beginIndex <= lastPrunedCacheBeginIndex) {
            return;
        }
        clusterIndexCache().keySet().removeIf(cacheIndex -> cacheIndex < beginIndex);
        lastPrunedCacheBeginIndex = beginIndex;
    }
}
