package compose.widgets

import androidx.compose.foundation.gestures.awaitEachGesture
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.padding
import androidx.compose.runtime.*
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.CacheDrawScope
import androidx.compose.ui.draw.clipToBounds
import androidx.compose.ui.draw.drawBehind
import androidx.compose.ui.draw.drawWithCache
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.geometry.Rect
import androidx.compose.ui.geometry.Size
import androidx.compose.ui.graphics.*
import androidx.compose.ui.graphics.drawscope.DrawStyle
import androidx.compose.ui.graphics.drawscope.Stroke
import androidx.compose.ui.input.pointer.PointerEventType
import androidx.compose.ui.input.pointer.pointerInput
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.text.AnnotatedString
import androidx.compose.ui.text.SpanStyle
import androidx.compose.ui.text.drawText
import androidx.compose.ui.text.rememberTextMeasurer
import androidx.compose.ui.unit.Density
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.dp
import androidx.compose.ui.unit.toSize
import compose.utils.contains
import kotlin.math.*


/**
 * Draws function line graphs.
 */
@Composable
fun FunctionLineGraph(
    modifier: Modifier,
    xRange: ClosedFloatingPointRange<Double>,
    yRange: ClosedFloatingPointRange<Double>,
    properties: GraphProperties = DefaultGraphProperties,
    lines: List<GraphLine>
) {
    val density = LocalDensity.current
    val styleAtXByLine = remember(density, lines) {
        lines.associateWith { it.lineStyleAtPoint.invoke(density) }
    }

    var drawLineModifier: Modifier = Modifier
    for (line in lines) {
        val styleAtByX = styleAtXByLine[line]!!
        drawLineModifier = drawLineModifier then remember(line, xRange, yRange, styleAtByX) {
            DrawLineModifier(
                xRange = xRange,
                yRange = yRange,
                line = line,
                styleAtX = styleAtByX
            )
        }
    }

    Box(modifier) {

        // Static content
        Box(Modifier
            .fillMaxSize()
            .graphicsLayer()
        ) {
            // Horizontal lines and their labels
            Spacer(Modifier
                .fillMaxSize()
                .padding(top = HoverXLabelAreaHeight, bottom = XAxisLabelsWidth)
                .drawHorizontalLines(yRange, properties)
            )

            // Vertical lines and their labels
            Spacer(Modifier
                .fillMaxSize()
                .padding(top = HoverXLabelAreaHeight, start = YAxisLabelsWidth)
                .drawVerticalLines(xRange, properties)
            )

            // The graph lines themselves
            Spacer(Modifier
                .fillMaxSize()
                .padding(top = HoverXLabelAreaHeight, bottom = XAxisLabelsWidth, start = YAxisLabelsWidth)
                .clipToBounds()
                .then(drawLineModifier)
            )
        }

        // Dynamic content
        Box(Modifier
            .fillMaxSize()
        ) {
            // The hover line
            Spacer(Modifier
                .fillMaxSize()
                .padding(top = HoverXLabelAreaHeight, bottom = XAxisLabelsWidth, start = YAxisLabelsWidth)
                .hoverLine(
                    xRange = xRange,
                    yRange = yRange,
                    properties = properties,
                    lines = lines,
                    styleAtXByLine = styleAtXByLine
                )
            )
        }
    }
}


/**
 * The width of the space for the labels on the Y axis.
 */
private val YAxisLabelsWidth = 80.dp


/**
 * The height of the space for the labels on the X axis.
 */
private val XAxisLabelsWidth = 20.dp


/**
 * The height of the area above the graph where the X value of the hover line is displayed.
 */
private val HoverXLabelAreaHeight = 20.dp


/**
 * Functions that help with calculations needed for drawing the graph and related objects.
 */
private class DrawHelpers(
    val xRange: ClosedFloatingPointRange<Double>,
    val yRange: ClosedFloatingPointRange<Double>,
    val line: GraphLine?,
    val size: Size
) {


    /**
     * The scale factor from pixels to function values on the X axis.
     */
    private val xFactor = xRange.size / size.width


    /**
     * The scale factor from pixels to function values on the Y axis.
     */
    private val yFactor = yRange.size / size.height


    /**
     * Converts from pixels to function values on the X axis.
     */
    fun xPxToFunction(xPx: Float) = xRange.start + xFactor * xPx


    /**
     * Converts from pixels to function values on the Y axis.
     */
    fun yPxToFunction(yPx: Float) = yRange.start + yFactor * (size.height - yPx)


    /**
     * Converts from function values to pixels on the X axis.
     */
    fun functionXToPx(x: Double) = (x - xRange.start) / xFactor


    /**
     * Returns the (pixel) value of the function at the given x location in pixels.
     */
    fun functionPx(xPx: Float): Float {
        if (line == null)
            throw IllegalStateException("Can't compute function; line has not been specified")

        return (size.height - (line.function(xPxToFunction(xPx)) - yRange.start) / yFactor).toFloat()
    }


}


/**
 * A [Modifier] that draws a single [GraphLine].
 */
fun DrawLineModifier(
    xRange: ClosedFloatingPointRange<Double>,
    yRange: ClosedFloatingPointRange<Double>,
    line: GraphLine,
    styleAtX: (x: Double, deltaX: Double) -> LineStyle?
) = Modifier.drawWithCache {
    with(DrawHelpers(xRange = xRange, yRange = yRange, size = size, line = line)) {
        val paths = mutableListOf<Path>()
        val styles = mutableListOf<LineStyle?>()

        val samplingRange = line.samplingRange ?: xRange
        val samplingPxStart = max(0.0, functionXToPx(samplingRange.start))
        val samplingPxEnd = min(size.width.toDouble(), functionXToPx(samplingRange.endInclusive))
        val samplingStepPx = line.samplingStepPx
        val samplingStepX = xPxToFunction(samplingStepPx)

        if (samplingPxStart > samplingPxEnd) {
            return@drawWithCache onDrawBehind {  }
        }

        var currentPath = Path().also { paths.add(it) }
        var currentStyle = styleAtX(xPxToFunction(0f), samplingStepX).also { styles.add(it) }

        currentPath.reset()
        currentPath.moveTo(0f, functionPx(0f))
        var x = samplingStepPx

        while (x <= samplingPxEnd) {
            val y = functionPx(x)
            currentPath.lineTo(x, y)
            x += samplingStepPx

            val style = styleAtX(xPxToFunction(x), samplingStepX)
            if (style != currentStyle) {
                currentPath = Path().also { paths.add(it) }
                currentStyle = style.also { styles.add(it) }
                currentPath.moveTo(x - samplingStepPx, y)
            }
        }

        onDrawBehind {
            for (i in 0..paths.lastIndex) {
                val style = styles[i] ?: continue
                drawPath(paths[i], brush = style.brush, style = style.drawStyle)
            }
        }
    }
}


/**
 * Computes the sequence of positions at which to draw horizontal/vertical lines.
 *
 * The first value of each pair is the pixel position; the 2nd one is the function value at that position.
 */
private fun CacheDrawScope.gridLinePositions(
    range: ClosedFloatingPointRange<Double>,
    canvasSize: Float,
    minDistance: Dp,
    invertedAxis: Boolean
): Sequence<Pair<Float, Double>> {
    val rangeSize = range.size
    val factor = canvasSize / rangeSize
    val stepScale = log10(rangeSize).roundToInt() - 1
    val step = 10.0.pow(stepScale).let { baseStep ->
        val stepMultiple = (1..10).firstOrNull { baseStep * it * factor >= minDistance.toPx() }
        baseStep * (stepMultiple ?: 10)
    }
    val stepPx = (step * factor).toFloat()
    if (stepPx == 0f)
        return emptySequence()

    val initialValue = ceil(range.start / step) * step
    val initialValuePx = ((initialValue - range.start) * factor).let {
        if (invertedAxis) canvasSize - it else it
    }.toFloat()

    return sequence {
        var value = initialValue
        var valuePx = initialValuePx

        val rangePx = 0 .. canvasSize.roundToInt()
        while (valuePx.roundToInt() in rangePx) {  // roundToInt prevents blinking of the top/last label
            yield(Pair(valuePx, value))
            valuePx += if (invertedAxis) -stepPx else stepPx
            value += step
        }
    }

}


/**
 * Draws the horizontal lines and values.
 */
@Composable
private fun Modifier.drawHorizontalLines(
    range: ClosedFloatingPointRange<Double>,
    props: GraphProperties,
): Modifier {
    val textMeasurer = rememberTextMeasurer(10)
    return drawWithCache {
        val positions = gridLinePositions(range, size.height, props.horizontalLinesMinDistance, invertedAxis = true)
        val offset = props.yLabelOffset.roundToPx()
        onDrawBehind {
            for ((yPx, y) in positions) {
                drawLine(props.lineColor, start = Offset(x = 0f, y = yPx), end = Offset(x = size.width, y = yPx))
                val text = props.yLabelFormatter(y)
                val textLayout = textMeasurer.measure(text, softWrap = false, maxLines = 1)
                val textTopLeft = Offset(x = 0f, y = yPx - offset - textLayout.size.height)
                val textBounds = Rect(offset = textTopLeft, size = textLayout.size.toSize())
                // Don't draw text that would be partially outside the canvas
                if (Rect(Offset.Zero, size).contains(textBounds))
                    drawText(textLayout, color = props.labelColor, topLeft = textTopLeft)
            }
        }
    }
}


/**
 * Draws the vertical lines and values.
 */
@Composable
private fun Modifier.drawVerticalLines(
    range: ClosedFloatingPointRange<Double>,
    props: GraphProperties,
): Modifier {
    val textMeasurer = rememberTextMeasurer(10)
    return drawWithCache {
        val positions = gridLinePositions(range, size.width, props.verticalLinesMinDistance, invertedAxis = false)
        val offset = props.yLabelOffset.roundToPx()
        onDrawBehind {
            for ((xPx, x) in positions) {
                drawLine(props.lineColor, start = Offset(x = xPx, y = 0f), end = Offset(x = xPx, y = size.height))
                val text = props.xLabelFormatter(x)
                val textLayout = textMeasurer.measure(text, softWrap = false, maxLines = 1)
                val textTopLeft = Offset(x = xPx + offset, y = size.height - textLayout.size.height)
                val textBounds = Rect(offset = textTopLeft, size = textLayout.size.toSize())
                // Don't draw text that would be partially outside the canvas
                if (Rect(Offset.Zero, size).contains(textBounds))
                    drawText(textLayout, color = props.labelColor, topLeft = textTopLeft)
            }
        }
    }
}


/**
 * Detects and draws the vertical line on hover.
 */
@Composable
private fun Modifier.hoverLine(
    xRange: ClosedFloatingPointRange<Double>,
    yRange: ClosedFloatingPointRange<Double>,
    properties: GraphProperties = DefaultGraphProperties,
    lines: List<GraphLine>,
    styleAtXByLine: Map<GraphLine, (x: Double, dx: Double) -> LineStyle?>
): Modifier {
    val textMeasurer = rememberTextMeasurer(10)
    var selectedX: Float? by remember { mutableStateOf(null) }
    return this
        .pointerInput(Unit) {
            awaitEachGesture {
                val event = awaitPointerEvent()
                when (event.type) {
                    PointerEventType.Enter,
                    PointerEventType.Move -> selectedX = event.changes.firstOrNull()?.position?.x
                    PointerEventType.Exit -> selectedX = null
                }
            }
        }
        .drawBehind {
            val shadowColor = properties.hoverLineProperties.textShadowColor
            val shadowStyle = if (shadowColor != null) {
                SpanStyle(
                    shadow = Shadow(
                        color = shadowColor,
                        offset = Offset(x = 1f, y = 1f)
                    )
                )
            } else
                SpanStyle()


            val xPx = selectedX ?: return@drawBehind
            val x = with(DrawHelpers(xRange = xRange, yRange = yRange, line = null, size = size)) {
                xPxToFunction(xPx)
            }

            // The text of the x value
            val xText = properties.hoverLineProperties.xFormatter(x)
            val xTextLayout = textMeasurer.measure(xText, softWrap = false, maxLines = 1)
            drawText(
                textLayoutResult = xTextLayout,
                color = properties.hoverLineProperties.color,
                topLeft = Offset(xPx - xTextLayout.size.width/2,  -xTextLayout.size.height * 1.2f)
            )

            // The vertical line itself
            drawLine(
                color = properties.hoverLineProperties.color,
                strokeWidth = properties.hoverLineProperties.width.toPx(),
                start = Offset(x = xPx, y = 0f),
                end = Offset(x = xPx, y = size.height)
            )


            // The circles marking the intersections, and the text near them
            for (line in lines) {
                val style = styleAtXByLine[line]!!.invoke(x, 0.0) ?: continue
                with(DrawHelpers(xRange = xRange, yRange = yRange, size = size, line = line)) {
                    val samplingRange = line.samplingRange ?: xRange
                    if (xPxToFunction(xPx) !in samplingRange)
                        return@with

                    val yPx = functionPx(xPx)
                    val y = yPxToFunction(yPx)

                    // The circle
                    drawCircle(
                        brush = style.brush,
                        center = Offset(x = xPx, y = yPx),
                        radius = properties.hoverLineProperties.circleRadius.toPx()
                    )

                    // The text of the function value at the intersection with the function
                    val yText = AnnotatedString(
                        text = properties.hoverLineProperties.yFormatter(line, x, y),
                        spanStyle = shadowStyle,
                    )
                    val yTextLayout = textMeasurer.measure(yText, softWrap = false, maxLines = 1)
                    val labelOffset = yTextLayout.size.height/5
                    drawText(
                        textLayoutResult = yTextLayout,
                        brush = style.brush,
                        topLeft = Offset(xPx + labelOffset, yPx - yTextLayout.size.height - labelOffset)
                    )
                }
            }
        }
}


/**
 * Defines a line in a [FunctionLineGraph].
 */
data class GraphLine(


    /**
     * The name of the line
     */
    val name: String,


    /**
     * The function to draw.
     */
    val function: (Double) -> Double,


    /**
     * The step, in pixels, for sampling the function.
     */
    val samplingStepPx: Float = 1f,


    /**
     * The range of this function to sample on the X axis. A `null` value indicates it should be sampled at all values
     * visible on the graph.
     */
    val samplingRange: ClosedFloatingPointRange<Double>? = null,


    /**
     * Returns a function that, given an x and a dx value for [function], returns the [LineStyle] with which the
     * function should be drawn between x and x+dx.
     * Additionally, it will be called with `dx=0.0` to determine the style of the circle at the intersection of the
     * vertical hover line and the graph line. The [LineStyle.brush] will be used for the text of the value at the
     * intersection.
     */
    val lineStyleAtPoint: Density.() -> ((x: Double, dx: Double) -> LineStyle?) = DefaultLineStyleAtPoint


) {


    init {
        require(samplingStepPx > 0) { "stepPx must be positive" }
    }


}


/**
 * Encapsulates the parameters of the style with which to draw a line.
 */
data class LineStyle(


    /**
     * The [Brush] with which to draw the line.
     */
    val brush: Brush,


    /**
     * The [DrawStyle] with which to draw the line.
     */
    val drawStyle: DrawStyle = Stroke(1f)


) {


    /**
     * Creates a [LineStyle] with which to draw the line.
     */
    constructor(


        /**
         * The color with which to draw the line.
         */
        color: Color,


        /**
         * The [DrawStyle] with which to draw the line.
         */
        drawStyle: DrawStyle = Stroke(1f)


    ): this(SolidColor(color), drawStyle)


}


/**
 * The default line drawing style.
 */
private val DefaultLineStyle = LineStyle(Color.Blue, Stroke(1f))


/**
 * Returns a function that returns the line style computed by the given function at all coordinates.
 */
inline fun FixedLineStyle(
    crossinline lineStyle: Density.() -> LineStyle
): Density.() -> ((Double, Double) -> LineStyle) {
    return {
        lineStyle().let {
            { _, _ -> it}
        }
    }
}


/**
 * Returns a function that returns the line style computed by the given function at all coordinates.
 */
fun FixedLineStyle(lineStyle: LineStyle): Density.() -> ((Double, Double) -> LineStyle) {
    return {
        { _, _ -> lineStyle }
    }
}


/**
 * The default draw style for [GraphLine].
 */
private val DefaultLineStyleAtPoint = FixedLineStyle(DefaultLineStyle)


/**
 * The size of the range.
 */
private val ClosedFloatingPointRange<Double>.size: Double
    get() = endInclusive - start


/**
 * Bundles the properties used in the drawing of the graph.
 */
data class GraphProperties(
    val lineColor: Color,
    val labelColor: Color,
    val yLabelOffset: Dp,
    val xLabelOffset: Dp,
    val horizontalLinesMinDistance: Dp,
    val verticalLinesMinDistance: Dp,
    val xLabelFormatter: (Double) -> String,
    val yLabelFormatter: (Double) -> String,
    val hoverLineProperties: GraphHoverLineProperties
)


/**
 * Bundles the properties of the vertical line displayed on mouse hover.
 */
data class GraphHoverLineProperties(
    val width: Dp,
    val color: Color,
    val circleRadius: Dp,
    val xFormatter: (Double) -> String,
    val yFormatter: (GraphLine, x: Double, y: Double) -> String,
    val textShadowColor: Color?,
)


/**
 * The default graph properties.
 */
val DefaultGraphProperties = GraphProperties(
    lineColor = Color.Black.copy(alpha = 0.2f),
    labelColor = Color.Black,
    yLabelOffset = 4.dp,
    xLabelOffset = 4.dp,
    horizontalLinesMinDistance = 50.dp,
    verticalLinesMinDistance = 80.dp,
    xLabelFormatter = Double::toString,
    yLabelFormatter = Double::toString,
    hoverLineProperties = GraphHoverLineProperties(
        width = Dp.Hairline,
        color = Color.Black,
        circleRadius = 4.dp,
        xFormatter = Double::toString,
        yFormatter = { _, _, y -> y.toString() },
        textShadowColor = Color.White
    )
)


