import type { RecoverableStatePersistenceOptions, Request } from '@crawlee/core';
import { RecoverableState } from '@crawlee/core';
import LogisticRegression from 'ml-logistic-regression';
import { Matrix } from 'ml-matrix';
import stringComparison from 'string-comparison';

export type RenderingType = 'clientOnly' | 'static';

type URLComponents = string[];

const urlComponents = (url: URL): URLComponents => {
    return [url.hostname, ...url.pathname.split('/')];
};

const calculateUrlSimilarity = (a: URLComponents, b: URLComponents): number | undefined => {
    const values: number[] = [];

    if (a[0] !== b[0]) {
        return 0;
    }

    for (let i = 1; i < Math.max(a.length, b.length); i++) {
        values.push(stringComparison.jaroWinkler.similarity(a[i] ?? '', b[i] ?? '') > 0.8 ? 1 : 0);
    }

    return sum(values) / Math.max(a.length, b.length);
};

const sum = (values: number[]) => values.reduce((acc, value) => acc + value);
const mean = (values: number[]) => (values.length > 0 ? sum(values) / values.length : undefined);

type FeatureVector = [staticResultsSimilarity: number, clientOnlyResultsSimilarity: number];

export interface RenderingTypePredictorOptions {
    /** A number between 0 and 1 that determines the desired ratio of rendering type detections */
    detectionRatio: number;
    persistenceOptions?: Partial<RecoverableStatePersistenceOptions>;
}

/**
 * Stores rendering type information for previously crawled URLs and predicts the rendering type for URLs that have yet to be crawled and recommends when rendering type detection should be performed.
 *
 * @experimental
 */
export class RenderingTypePredictor {
    private detectionRatio: number;
    private state: RecoverableState<{
        logreg: LogisticRegression;
        detectionResults: Map<RenderingType, Map<string | undefined, URLComponents[]>>;
    }>;

    constructor({ detectionRatio, persistenceOptions }: RenderingTypePredictorOptions) {
        this.detectionRatio = detectionRatio;
        this.state = new RecoverableState({
            defaultState: {
                logreg: new LogisticRegression({ numSteps: 1000, learningRate: 0.05 }),
                detectionResults: new Map<RenderingType, Map<string | undefined, URLComponents[]>>(),
            },
            serialize: (state) =>
                JSON.stringify({
                    logreg: state.logreg.toJSON(),
                    detectionResults: Array.from(state.detectionResults.entries()).map(
                        ([renderingType, urlPartsByLabel]) => ({
                            renderingType,
                            urlPartsByLabel: Array.from(urlPartsByLabel.entries()).map(([label, urlParts]) => ({
                                label,
                                urlParts,
                            })),
                        }),
                    ),
                }),
            deserialize: (serializedState) => {
                const { logreg, detectionResults = [] } = JSON.parse(serializedState);

                return {
                    logreg: LogisticRegression.load(logreg),
                    detectionResults: new Map(
                        detectionResults.map((serializedItem: any) => [
                            serializedItem.renderingType,
                            new Map(serializedItem.urlPartsByLabel.map((item: any) => [item.label, item.urlParts])),
                        ]),
                    ),
                };
            },
            persistStateKey: 'rendering-type-predictor-state',
            persistenceEnabled: true,
            ...persistenceOptions,
        });
    }

    /**
     * Initialize the predictor by restoring persisted state.
     */
    async initialize(): Promise<void> {
        await this.state.initialize();
    }

    /**
     * Predict the rendering type for a given URL and request label.
     */
    public predict({ url, loadedUrl, label }: Request): {
        renderingType: RenderingType;
        detectionProbabilityRecommendation: number;
    } {
        const { logreg } = this.state.currentValue;
        if (logreg.classifiers.length === 0) {
            return { renderingType: 'clientOnly', detectionProbabilityRecommendation: 1 };
        }

        const predictionUrl = new URL(loadedUrl ?? url);

        const urlFeature = new Matrix([this.calculateFeatureVector(urlComponents(predictionUrl), label)]);
        const [prediction] = logreg.predict(urlFeature);
        const scores = [logreg.classifiers[0].testScores(urlFeature), logreg.classifiers[1].testScores(urlFeature)];

        return {
            renderingType: prediction === 1 ? 'static' : 'clientOnly',
            detectionProbabilityRecommendation:
                Math.abs(scores[0] - scores[1]) < 0.1
                    ? 1
                    : this.detectionRatio * Math.max(1, 5 - this.resultCount(label)),
        };
    }

    /**
     * Store the rendering type for a given URL and request label. This updates the underlying prediction model, which may be costly.
     */
    public storeResult(requests: Request | Request[], renderingType: RenderingType) {
        const state = this.state.currentValue;

        for (const { url, loadedUrl, label } of Array.isArray(requests) ? requests : [requests]) {
            const resultUrl = new URL(loadedUrl ?? url);

            if (!state.detectionResults.has(renderingType)) {
                state.detectionResults.set(renderingType, new Map());
            }

            if (!state.detectionResults.get(renderingType)!.has(label)) {
                state.detectionResults.get(renderingType)!.set(label, []);
            }

            state.detectionResults.get(renderingType)!.get(label)!.push(urlComponents(resultUrl));
        }

        this.retrain();
    }

    private resultCount(label: string | undefined): number {
        return Array.from(this.state.currentValue.detectionResults.values())
            .map((results) => results.get(label)?.length ?? 0)
            .reduce((acc, value) => acc + value, 0);
    }

    protected calculateFeatureVector(url: URLComponents, label: string | undefined): FeatureVector {
        return [
            mean(
                (this.state.currentValue.detectionResults.get('static')?.get(label) ?? []).map(
                    (otherUrl) => calculateUrlSimilarity(url, otherUrl) ?? 0,
                ),
            ) ?? 0,
            mean(
                (this.state.currentValue.detectionResults.get('clientOnly')?.get(label) ?? []).map(
                    (otherUrl) => calculateUrlSimilarity(url, otherUrl) ?? 0,
                ),
            ) ?? 0,
        ];
    }

    protected retrain(): void {
        const X: FeatureVector[] = [
            [0, 1],
            [1, 0],
        ];
        const Y: number[] = [0, 1];

        for (const [renderingType, urlsByLabel] of this.state.currentValue.detectionResults.entries()) {
            for (const [label, urls] of urlsByLabel) {
                for (const url of urls) {
                    X.push(this.calculateFeatureVector(url, label));
                    Y.push(renderingType === 'static' ? 1 : 0);
                }
            }
        }

        this.state.currentValue.logreg.train(new Matrix(X), Matrix.columnVector(Y));
    }
}
