const lcjs = require('@lightningchart/lcjs')

const {
    lightningChart,
    AxisScrollStrategies,
    AxisTickStrategies,
    ColorRGBA,
    PalettedFill,
    LUT,
    SolidFill,
    SolidLine,
    emptyLine,
    emptyFill,
    regularColorSteps,
} = lcjs

const EXAMPLE_PATH = 'examples/assets/1719/'
const HISTORY_SECONDS = 15
const HISTORY_MS = HISTORY_SECONDS * 1000
const MAX_FRAME_CATCH_UP_SECONDS = 0.12
const PAUSE_RESET_SECONDS = 0.5
const AREA_TARGET_POINTS = 96
const AREA_REFRESH_HZ = 24
const AREA_DECAY_STEPS = 12
const JAMMING_TARGET_MAX_COUNT = 12
const JAMMING_TARGET_SMOOTH_RADIUS = 3
const JAMMING_TARGET_THRESHOLD_ABOVE_FLOOR = 0.06
const JAMMING_TARGET_MIN_POWER = 0.12
const JAMMING_TARGET_MIN_BANDWIDTH_MHZ = 0.3
const JAMMING_TARGET_MERGE_GAP_MHZ = 1.4
const JAMMING_TARGET_ACTIVITY_WINDOW_SECONDS = 1.8
const JAMMING_TARGET_REFRESH_HZ = 8
const JAMMING_TARGET_TRACK_SMOOTHING = 0.12
const JAMMING_TARGET_TRACK_HOLD_SECONDS = 1.6
const JAMMING_TARGET_TRACK_MATCH_GAP_MHZ = 2.0

const exampleContainer = document.getElementById('chart') || document.body
if (exampleContainer === document.body) {
    exampleContainer.style.width = '100vw'
    exampleContainer.style.height = '100vh'
    exampleContainer.style.margin = '0px'
}
exampleContainer.style.display = 'flex'
exampleContainer.style.flexDirection = 'column'

const chartContainer = document.createElement('div')
exampleContainer.append(chartContainer)
chartContainer.style.width = '100%'
chartContainer.style.height = '100%'

const assetUrl = (fileName) => document.head.baseURI + EXAMPLE_PATH + fileName
const formatMHz = (hz) => `${(hz / 1_000_000).toFixed(3)}`

const loadArrayBuffer = async (url) => {
    const response = await fetch(url)
    if (!response.ok) throw new Error(`Failed to load ${url}: ${response.status}`)
    return response.arrayBuffer()
}

const sliceRow = (array, rowIndex, binCount) => array.subarray(rowIndex * binCount, (rowIndex + 1) * binCount)

const makeAreaSpectrum = (row, minMHz, stepMHz, binAggregation, valueScale) => {
    const pointCount = Math.ceil(row.length / binAggregation)
    const xValues = new Float64Array(pointCount)
    const yValues = new Float32Array(pointCount)
    for (let i = 0; i < pointCount; i += 1) {
        const iStart = i * binAggregation
        const iEnd = Math.min(row.length, iStart + binAggregation)
        let sum = 0
        for (let j = iStart; j < iEnd; j += 1) {
            sum += row[j] * valueScale
        }
        xValues[i] = minMHz + ((iStart + iEnd - 1) / 2) * stepMHz
        yValues[i] = sum / (iEnd - iStart)
    }
    return { xValues, yValues }
}

const makeTemporalActivitySpectrum = (spectrogram, endRowIndex, rowCount, binCount, windowRows, output, valueScale) => {
    output.fill(0)
    for (let i = 0; i < windowRows; i += 1) {
        const sourceRow = (endRowIndex - i + rowCount) % rowCount
        const row = sliceRow(spectrogram, sourceRow, binCount)
        for (let j = 0; j < binCount; j += 1) {
            output[j] += row[j] * valueScale
        }
    }
    const scale = 1 / windowRows
    for (let i = 0; i < binCount; i += 1) {
        output[i] *= scale
    }
    return output
}

const estimateNoiseFloor = (row) => {
    const sampleCount = Math.min(256, row.length)
    const stride = Math.max(1, Math.floor(row.length / sampleCount))
    const samples = []
    for (let i = 0; i < row.length && samples.length < sampleCount; i += stride) {
        samples.push(row[i])
    }
    samples.sort((a, b) => a - b)
    const quietSampleCount = Math.max(1, Math.floor(samples.length * 0.55))
    let sum = 0
    for (let i = 0; i < quietSampleCount; i += 1) {
        sum += samples[i]
    }
    return sum / quietSampleCount
}

const mergeNearbyJammingTargets = (targets) => {
    const sortedTargets = targets.slice().sort((a, b) => a.freqStart - b.freqStart)
    const mergedTargets = []
    for (const target of sortedTargets) {
        const previous = mergedTargets[mergedTargets.length - 1]
        if (previous && target.freqStart - previous.freqEnd <= JAMMING_TARGET_MERGE_GAP_MHZ) {
            const previousBandwidth = previous.freqEnd - previous.freqStart
            const targetBandwidth = target.freqEnd - target.freqStart
            const totalBandwidth = previousBandwidth + targetBandwidth
            previous.freqEnd = target.freqEnd
            previous.avgAmplitude =
                totalBandwidth > 0
                    ? (previous.avgAmplitude * previousBandwidth + target.avgAmplitude * targetBandwidth) / totalBandwidth
                    : Math.max(previous.avgAmplitude, target.avgAmplitude)
            previous.peakAmplitude = Math.max(previous.peakAmplitude, target.peakAmplitude)
            previous.lastSeen = Math.max(previous.lastSeen ?? 0, target.lastSeen ?? 0)
        } else {
            mergedTargets.push({ ...target })
        }
    }
    return mergedTargets
}

const getTargetGapMHz = (a, b) => Math.max(0, Math.max(a.freqStart, b.freqStart) - Math.min(a.freqEnd, b.freqEnd))

const smoothJammingTargets = (previousTargets, detectedTargets, now) => {
    const nextTargets = []
    const usedDetected = new Set()
    for (const previous of previousTargets) {
        let bestMatchIndex = -1
        let bestMatchGap = Number.POSITIVE_INFINITY
        detectedTargets.forEach((target, i) => {
            if (usedDetected.has(i)) return
            const gap = getTargetGapMHz(previous, target)
            if (gap <= JAMMING_TARGET_TRACK_MATCH_GAP_MHZ && gap < bestMatchGap) {
                bestMatchIndex = i
                bestMatchGap = gap
            }
        })

        if (bestMatchIndex >= 0) {
            const target = detectedTargets[bestMatchIndex]
            usedDetected.add(bestMatchIndex)
            nextTargets.push({
                freqStart: previous.freqStart * (1 - JAMMING_TARGET_TRACK_SMOOTHING) + target.freqStart * JAMMING_TARGET_TRACK_SMOOTHING,
                freqEnd: previous.freqEnd * (1 - JAMMING_TARGET_TRACK_SMOOTHING) + target.freqEnd * JAMMING_TARGET_TRACK_SMOOTHING,
                avgAmplitude:
                    previous.avgAmplitude * (1 - JAMMING_TARGET_TRACK_SMOOTHING) + target.avgAmplitude * JAMMING_TARGET_TRACK_SMOOTHING,
                peakAmplitude: Math.max(previous.peakAmplitude, target.peakAmplitude),
                lastSeen: now,
            })
        } else if (now - previous.lastSeen <= JAMMING_TARGET_TRACK_HOLD_SECONDS * 1000) {
            nextTargets.push(previous)
        }
    }

    detectedTargets.forEach((target, i) => {
        if (!usedDetected.has(i)) {
            nextTargets.push({ ...target, lastSeen: now })
        }
    })

    return mergeNearbyJammingTargets(nextTargets)
        .sort((a, b) => b.avgAmplitude - a.avgAmplitude)
        .slice(0, JAMMING_TARGET_MAX_COUNT)
        .sort((a, b) => a.freqStart - b.freqStart)
}

const detectJammingTargets = (row, minMHz, stepMHz) => {
    const noiseFloor = estimateNoiseFloor(row)
    const activeThreshold = Math.max(JAMMING_TARGET_MIN_POWER, noiseFloor + JAMMING_TARGET_THRESHOLD_ABOVE_FLOOR)
    const minBins = Math.max(3, Math.round(JAMMING_TARGET_MIN_BANDWIDTH_MHZ / stepMHz))
    const maxGapBins = Math.max(1, Math.round(JAMMING_TARGET_MERGE_GAP_MHZ / stepMHz))
    const targets = []
    let currentTarget

    const closeTarget = () => {
        if (!currentTarget) return
        const binWidth = currentTarget.endBin - currentTarget.startBin + 1
        if (binWidth >= minBins && currentTarget.activeCount > 0) {
            targets.push({
                freqStart: minMHz + currentTarget.startBin * stepMHz,
                freqEnd: minMHz + currentTarget.endBin * stepMHz,
                avgAmplitude: currentTarget.amplitudeSum / currentTarget.activeCount,
                peakAmplitude: currentTarget.peakAmplitude,
            })
        }
        currentTarget = undefined
    }

    for (let i = 0; i < row.length; i += 1) {
        const iStart = Math.max(0, i - JAMMING_TARGET_SMOOTH_RADIUS)
        const iEnd = Math.min(row.length - 1, i + JAMMING_TARGET_SMOOTH_RADIUS)
        let smoothSum = 0
        for (let j = iStart; j <= iEnd; j += 1) {
            smoothSum += row[j]
        }
        const amplitude = smoothSum / (iEnd - iStart + 1)
        const active = amplitude >= activeThreshold

        if (active) {
            if (!currentTarget) {
                currentTarget = {
                    startBin: i,
                    endBin: i,
                    activeCount: 0,
                    amplitudeSum: 0,
                    peakAmplitude: 0,
                    gapBins: 0,
                }
            }
            currentTarget.endBin = i
            currentTarget.activeCount += 1
            currentTarget.amplitudeSum += amplitude
            currentTarget.peakAmplitude = Math.max(currentTarget.peakAmplitude, amplitude)
            currentTarget.gapBins = 0
        } else if (currentTarget) {
            currentTarget.gapBins += 1
            if (currentTarget.gapBins > maxGapBins) {
                closeTarget()
            }
        }
    }
    closeTarget()

    return mergeNearbyJammingTargets(targets)
        .sort((a, b) => b.avgAmplitude - a.avgAmplitude)
        .slice(0, JAMMING_TARGET_MAX_COUNT)
        .sort((a, b) => a.freqStart - b.freqStart)
}

const run = async () => {
    const metadata = await fetch(assetUrl('metadata.json')).then((response) => response.json())
    const spectrogramBuffer = await loadArrayBuffer(assetUrl(metadata.spectrogram.file))

    const binCount = metadata.binCount
    const rowCount = metadata.rowCount
    const rowRate = metadata.rowRate
    const rowStepMs = 1000 / rowRate
    const frequencyMinMHz = metadata.frequencyAxis.minHz / 1_000_000
    const frequencyMaxMHz = metadata.frequencyAxis.maxHz / 1_000_000
    const frequencyStepMHz = (frequencyMaxMHz - frequencyMinMHz) / (binCount - 1)
    const areaBinAggregation = Math.max(1, Math.round(binCount / AREA_TARGET_POINTS))
    const isUint8Spectrogram = metadata.spectrogram.format === 'uint8'
    const spectrogram = isUint8Spectrogram ? new Uint8Array(spectrogramBuffer) : new Float32Array(spectrogramBuffer)
    const valueScale = metadata.spectrogram.normalizedScale ?? (isUint8Spectrogram ? 1 / 255 : 1)
    const paletteMax = isUint8Spectrogram ? 255 : 1
    const expectedValues = rowCount * binCount

    if (spectrogram.length !== expectedValues) {
        throw new Error(`RF asset shape mismatch. Expected ${expectedValues} values, got ${spectrogram.length}.`)
    }

    const lc = lightningChart()
    const chart = lc
        .ChartXY({
            container: chartContainer,
            defaultAxisX: { type: 'linear-highPrecision' },
            defaultAxisY: { type: 'linear-highPrecision' },
            legend: { visible: false },
        })
        .setTitle(`RF Jamming Target Identification | ${frequencyMinMHz.toFixed(3)}-${frequencyMaxMHz.toFixed(3)} MHz`)
        .setCursorMode('show-nearest')
        .setPadding({ left: 24 })

    chart.yAxis.dispose()
    const axisSpectrum = chart
        .addAxisY({ iStack: 2 })
        .setTitle('Normalized power')
        .setInterval({ start: 0, end: 1 })
        .setLength({ pixels: 180 })
        .setMargins(10, 0)
        .setUserInteractions({
            zoom: false,
            pan: false,
            rectangleZoom: false,
            restoreDefault: false,
            chartInteractions: { zoom: false, pan: false, rectangleZoom: false, restoreDefault: false },
        })
    const axisJammingTargets = chart
        .addAxisY({ iStack: 1 })
        .setTitle('Jamming targets')
        .setInterval({ start: 0, end: 1 })
        .setLength({ pixels: 120 })
        .setMargins(6, 6)
        .setUserInteractions({
            zoom: false,
            pan: false,
            rectangleZoom: false,
            restoreDefault: false,
            chartInteractions: { zoom: false, pan: false, rectangleZoom: false, restoreDefault: false },
        })
    const axisSpectrogram = chart
        .addAxisY({ iStack: 0, type: 'linear-highPrecision' })
        .setScrollStrategy(AxisScrollStrategies.scrolling)
        .setTickStrategy(AxisTickStrategies.DateTime, (s) =>
            s.setCursorFormatter((value, range, locale) =>
                new Date(value).toLocaleTimeString(locale, {
                    minute: '2-digit',
                    hour: '2-digit',
                    second: '2-digit',
                    hour12: false,
                }),
            ),
        )
        .setDefaultInterval((state) => ({
            end: state.dataMax ?? 0,
            start: (state.dataMax ?? 0) - HISTORY_MS,
            stopAxisAfter: false,
        }))

    chart.axisX
        .setTitle('RF frequency')
        .setUnits('MHz')
        .setInterval({ start: frequencyMinMHz, end: frequencyMaxMHz })
        .setUserInteractions({
            zoom: {
                mode: 'toward-pointer',
            },
        })
        .setScrollMargins(false)

    const areaSeriesList = new Array(AREA_DECAY_STEPS).fill(0).map((_, i) =>
        chart
            .addAreaSeries({ yAxis: axisSpectrum })
            .setCurvePreprocessing({ type: 'spline' })
            .setHighlightOnHover(false)
            .setName(i === AREA_DECAY_STEPS - 1 ? 'Latest activity' : 'Activity history')
            .setCursorEnabled(i === AREA_DECAY_STEPS - 1 ? true : false),
    )
    const defaultAreaColor = areaSeriesList[AREA_DECAY_STEPS - 1].getStrokeStyle().getFillStyle().getColor()
    const areaStyles = new Array(AREA_DECAY_STEPS).fill(0).map((_, i) => {
        const age = AREA_DECAY_STEPS - 1 - i
        const fade = 1 - age / AREA_DECAY_STEPS
        const strokeAlpha = Math.round(210 * fade)
        const fillAlpha = Math.round(120 * fade)
        const strokeColor = ColorRGBA(defaultAreaColor.getR(), defaultAreaColor.getG(), defaultAreaColor.getB(), strokeAlpha)
        const fillColor = ColorRGBA(defaultAreaColor.getR(), defaultAreaColor.getG(), defaultAreaColor.getB(), fillAlpha)
        return {
            stroke: new SolidLine({ thickness: 1, fillStyle: new SolidFill({ color: strokeColor }) }),
            fill: new SolidFill({ color: fillColor }),
        }
    })
    areaSeriesList.forEach((series, i) => {
        const style = areaStyles[i]
        series
            .setStrokeStyle(emptyLine)
            // .setStrokeStyle(style.stroke)
            .setAreaFillStyle(style.fill)
    })

    const theme = chart.getTheme()
    const isDarkTheme = theme.isDark
    const jammingTargetStrokeStyle = new SolidLine({
        thickness: 2,
        fillStyle: new SolidFill({
            color: isDarkTheme ? ColorRGBA(255, 80, 62, 235) : ColorRGBA(180, 34, 24, 245),
        }),
    })
    const jammingTargetSeries = chart
        .addRectangleSeries({ yAxis: axisJammingTargets })
        .setName('Identified jamming ranges')
        .setCursorEnabled(true)
        .setDefaultStyle((figure) => figure.setFillStyle(emptyFill).setStrokeStyle(jammingTargetStrokeStyle))
    jammingTargetSeries.setCursorBehavior({
        location: (info) => ({ x: (info.x1 + info.x2) / 2, y: info.y2 }),
    })

    const palette = new PalettedFill({
        lookUpProperty: 'value',
        lut: new LUT({
            interpolate: true,
            steps: regularColorSteps(0, paletteMax, theme.examples.spectrogramColorPalette),
            units: 'power',
        }),
    })

    const heatmapSeries = chart
        .addHeatmapScrollingGridSeries({
            scrollDimension: 'rows',
            resolution: binCount,
            yAxis: axisSpectrogram,
        })
        .setName('Power')
        .setStep({ x: frequencyStepMHz, y: rowStepMs })
        .setFillStyle(palette)
        .setWireframeStyle(emptyLine)
        .setIntensityInterpolation('disabled')
        .setAggregation('max')
        .setDataCleaning({ minDataPointCount: Math.ceil(HISTORY_SECONDS * rowRate * 1.5) })

    let rowIndex = 0
    let tPrevious = performance.now()
    let rowCarry = 0
    let valuesPushed = 0
    let tStats = performance.now()
    let tLastArea = 0
    let tLastJammingTargets = 0
    let tFirstRow
    let nextHeatmapSampleIndex = 0
    let detectedTargetCount = 0
    const areaHistory = []
    let stableJammingTargets = []
    const jammingActivityWindowRows = Math.max(1, Math.round(rowRate * JAMMING_TARGET_ACTIVITY_WINDOW_SECONDS))
    const jammingActivitySpectrum = new Float32Array(binCount)
    const baseTitle = chart.getTitle()

    document.addEventListener('visibilitychange', () => {
        if (!document.hidden) {
            tPrevious = performance.now()
            rowCarry = 0
        }
    })

    const collectRows = (count) => {
        const rows = []
        for (let i = 0; i < count; i += 1) {
            rows.push(sliceRow(spectrogram, (rowIndex + i) % rowCount, binCount))
        }
        return rows
    }

    const appendRows = (count) => {
        if (count <= 0) return
        if (!tFirstRow) {
            tFirstRow = Date.now()
            heatmapSeries.setStart({ x: frequencyMinMHz, y: tFirstRow })
        }
        const rows = collectRows(count)
        heatmapSeries.invalidateIntensityValues({
            iSample: nextHeatmapSampleIndex,
            values: rows,
        })
        nextHeatmapSampleIndex += count

        const latestSourceRow = (rowIndex + count - 1) % rowCount
        if (performance.now() - tLastArea >= 1000 / AREA_REFRESH_HZ) {
            const latestRow = sliceRow(spectrogram, latestSourceRow, binCount)
            areaHistory.unshift(makeAreaSpectrum(latestRow, frequencyMinMHz, frequencyStepMHz, areaBinAggregation, valueScale))
            areaHistory.length = Math.min(areaHistory.length, AREA_DECAY_STEPS)
            areaSeriesList.forEach((series, i) => {
                const historyIndex = AREA_DECAY_STEPS - 1 - i
                const area = areaHistory[historyIndex]
                series.clear()
                if (area) series.appendSamples(area)
            })
            tLastArea = performance.now()
        }
        if (performance.now() - tLastJammingTargets >= 1000 / JAMMING_TARGET_REFRESH_HZ) {
            const activityRow = makeTemporalActivitySpectrum(
                spectrogram,
                latestSourceRow,
                rowCount,
                binCount,
                Math.min(jammingActivityWindowRows, rowCount),
                jammingActivitySpectrum,
                valueScale,
            )
            const jammingTargets = detectJammingTargets(activityRow, frequencyMinMHz, frequencyStepMHz)
            stableJammingTargets = smoothJammingTargets(stableJammingTargets, jammingTargets, performance.now())
            jammingTargetSeries.clear()
            stableJammingTargets.forEach((target) => {
                jammingTargetSeries.add({
                    x1: target.freqStart,
                    x2: target.freqEnd,
                    y1: 0,
                    y2: target.avgAmplitude,
                })
            })
            detectedTargetCount = stableJammingTargets.length
            tLastJammingTargets = performance.now()
        }

        rowIndex += count
        valuesPushed += count * binCount
    }

    const animate = () => {
        const now = performance.now()
        const rawElapsedSeconds = (now - tPrevious) / 1000
        const wasPaused = rawElapsedSeconds > PAUSE_RESET_SECONDS
        if (wasPaused) {
            rowCarry = 0
        }
        const elapsedSeconds = wasPaused ? MAX_FRAME_CATCH_UP_SECONDS : Math.min(rawElapsedSeconds, MAX_FRAME_CATCH_UP_SECONDS)
        let rowsToAppend = elapsedSeconds * rowRate + rowCarry
        rowCarry = rowsToAppend % 1
        rowsToAppend = Math.floor(rowsToAppend)
        tPrevious = now

        appendRows(rowsToAppend)

        if (now - tStats >= 1000) {
            const pointsPerSecond = Math.round((valuesPushed * 1000) / (now - tStats))
            chart.setTitle(`${baseTitle} | ${detectedTargetCount} targets | ${pointsPerSecond.toLocaleString()} values/s`)
            valuesPushed = 0
            tStats = now
        }

        requestAnimationFrame(animate)
    }

    const initialRows = Math.min(rowCount, Math.ceil(HISTORY_SECONDS * rowRate))
    tFirstRow = Date.now() - (initialRows - 1) * rowStepMs
    heatmapSeries.setStart({ x: frequencyMinMHz, y: tFirstRow })
    appendRows(initialRows)
    tPrevious = performance.now()
    requestAnimationFrame(animate)
}

run().catch((error) => {
    console.error(error)
    const message = document.createElement('pre')
    message.textContent = error && error.stack ? error.stack : String(error)
    message.style.padding = '16px'
    message.style.margin = '0px'
    exampleContainer.append(message)
})