class Points extends THREE.Object3D {
    constructor(
        data,
        r = 1.0,
        g = 1.0,
        b = 1.0,
        sprite = undefined,
        particleSize = 5,
        z = 0
    ) {
        // :param data: points at a particular flow layer with shape (nDimensions, nSamples).
        super();

        this.data = data
        this.dim0 = null;
        this.dim1 = null;
        this.points = null;
        this.colors = null;
        this.z = z;

        this.geometry = new THREE.BufferGeometry();
        this.material = new THREE.PointsMaterial({
            size: particleSize,
            vertexColors: THREE.VertexColors,
            map: sprite,
            alphaTest: 0.5,
            transparent: true
        });
        this.points = new THREE.Points(this.geometry, this.material);

        this.add(this.points);  // Make this.points a child of this class

        this.setDimensions(0, 1)
        this.setColors(r, g, b);
    }

    xs() {
        return this.data[this.dim0];
    }

    ys() {
        return this.data[this.dim1];
    }

    setDimensions(dim0, dim1) {
        this.dim0 = dim0;
        this.dim1 = dim1;
        this.setGeometry();
    }

    setColors(r, g, b) {
        this.colors = [];
        for (let i = 0; i < this.data[this.dim0].length; i++) {
            let color = new THREE.Color();
            color.setRGB(r, g, b);
            this.colors.push(color.r, color.g, color.b);
        }
        this.geometry.setAttribute('color', new THREE.Float32BufferAttribute(this.colors, 3));
    }

    setGeometry() {
        let xs = this.xs();
        let ys = this.ys();
        let z = this.z;

        let positions = [];
        for (let i = 0; i < xs.length; i++) {
            positions.push(xs[i], ys[i], z);
        }
        this.geometry.setAttribute('position', new THREE.Float32BufferAttribute(positions, 3));
    }
}

class PathLines extends THREE.Object3D {
    constructor(data, z = -1, linewidth = 1) {
        // :param data: points with shape (nDimensions, nSamples, nLayers).
        super();

        this.dim0 = 0
        this.dim1 = 1
        this.z = z

        let nDimensions = data.length;
        let nSamples = data[0].length;
        let nLayers = data[0][0].length;  // Actually one more than the number of layers


        this.sampleArray = [];  // shape = (nSamples, nDimensions, nLayers)
        for (let sampleIdx = 0; sampleIdx < nSamples; sampleIdx++) {
            let dimensionArray = [];
            for (let dimIdx = 0; dimIdx < nDimensions; dimIdx++) {
                let layerArray = [];
                for (let layerIdx = 0; layerIdx < nLayers; layerIdx++) {
                    layerArray.push(data[dimIdx][sampleIdx][layerIdx])
                }  // nLayers
                dimensionArray.push(layerArray)
            }  // nDimensions
            this.sampleArray.push(dimensionArray);
        }  // nSamples

        this.geometry = new THREE.BufferGeometry();
        this.material = new THREE.LineBasicMaterial({
            color: 0xbfbfbf,
            linewidth: linewidth
        });

        this.indices = [];
        for (let sampleIdx = 0; sampleIdx < nSamples; sampleIdx++) {
            for (let layerIdx = 0; layerIdx < nLayers; layerIdx++) {
                if (layerIdx < nLayers - 1) {
                    this.indices.push(sampleIdx * nLayers + layerIdx)
                    this.indices.push(sampleIdx * nLayers + layerIdx + 1)
                }
            }
        }
        this.geometry.setIndex(new THREE.BufferAttribute(new Uint16Array(this.indices), 1));
        this.setGeometry()

        this.lines = new THREE.LineSegments(this.geometry, this.material)
        this.add(this.lines)
    }

    setDimensions(dim0, dim1) {
        this.dim0 = dim0
        this.dim1 = dim1
        this.setGeometry()
    }

    setGeometry() {
        let vertices = [];
        let nLayers = this.sampleArray[0][this.dim0].length;  // Actually one more than the number of layers
        let nSamples = this.sampleArray.length;
        for (let sampleIdx = 0; sampleIdx < nSamples; sampleIdx++) {
            for (let layerIdx = 0; layerIdx < nLayers; layerIdx++) {
                vertices.push(this.sampleArray[sampleIdx][this.dim0][layerIdx])
                vertices.push(this.sampleArray[sampleIdx][this.dim1][layerIdx])
                vertices.push(this.z)
            }
        }

        this.geometry.setAttribute('position', new THREE.BufferAttribute(new Float32Array(vertices), 3));
    }
}

class PathLayer extends THREE.Object3D {
    constructor(
        data,
        layerIdx,
        r = 1.0,
        g = 1.0,
        b = 1.0,
        sprite = undefined,
        particleSize = 5,
        z = 0
    ) {
        // :param data: points with shape (nDimensions, nSamples, nLayers).
        super();

        let nDimensions = data.length;
        let nSamples = data[0].length;

        let pointsArray = [];
        for (let dimIdx = 0; dimIdx < nDimensions; dimIdx++) {
            let dimensionArray = [];
            for (let sampleIdx = 0; sampleIdx < nSamples; sampleIdx++) {
                dimensionArray.push(data[dimIdx][sampleIdx][layerIdx]);
            }  // nDimensions
            pointsArray.push(dimensionArray);
        }  // nSamples

        this.points = new Points(pointsArray, r, g, b, sprite, particleSize, z)
        this.add(this.points)
    }

    setDimensions(dim0, dim1) {
        this.points.setDimensions(dim0, dim1)
    }
}

class PathsDataset extends THREE.Object3D {
    constructor(
        data,
        rDataLayer = 1.0,
        gDataLayer = 1.0,
        bDataLayer = 1.0,
        spriteDataLayer = undefined,
        rLatentLayer = 1.0,
        gLatentLayer = 1.0,
        bLatentLayer = 1.0,
        spriteLatentLayer = undefined,
        particleSize = 5,
        z = 0
    ) {
        // :param data: points with shape (nDimensions, nSamples, nLayers).
        super();

        let nLayers = data[0][0].length;

        this.dataLayer = new PathLayer(
            data,
            0,
            rDataLayer,
            gDataLayer,
            bDataLayer,
            spriteDataLayer,
            particleSize,
            z
        );
        this.latentLayer = new PathLayer(
            data,
            nLayers - 1,
            rLatentLayer,
            gLatentLayer,
            bLatentLayer,
            spriteLatentLayer,
            particleSize,
            z
        );
        this.paths = new PathLines(data, z - 1)

        this.add(this.dataLayer);
        this.add(this.latentLayer);
        this.add(this.paths);
    }

    setDimensions(dim0, dim1) {
        this.dataLayer.setDimensions(dim0, dim1);
        this.latentLayer.setDimensions(dim0, dim1);
        this.paths.setDimensions(dim0, dim1);
    }
}

class CovarianceEllipse extends THREE.Object3D {
    constructor(sigma = 1, nPoints = 100, linewidth = 3) {
        super();

        this.curve = new THREE.EllipseCurve(
            0, 0,            // ax, aY
            sigma, sigma,           // xRadius, yRadius
            0, 2 * Math.PI,  // aStartAngle, aEndAngle
            false,            // aClockwise
            0                 // aRotation
        );

        this.geometry = new THREE.BufferGeometry().setFromPoints(this.curve.getPoints(nPoints));
        this.material = new THREE.LineBasicMaterial({color: 0x00ff00, linewidth: linewidth});

        // Create the final object to add to the scene
        this.ellipse = new THREE.Line(this.geometry, this.material);
        this.add(this.ellipse)
    }
}

function setCameraParameters(dim0, dim1, limits, limitScaling, camera) {
    // "Axis" limits in data space. If our points are drawn from a 2D Uniform(-1, 1),
    // then having limits (-1.05, 1.05) for both axes will let us see all points.
    let bounds = computeRenderBounds(dim0, dim1, limits, limitScaling);

    camera.left = bounds.xMin;
    camera.right = bounds.xMax;
    camera.top = bounds.yMax;
    camera.bottom = bounds.yMin;

    camera.updateProjectionMatrix();
}

function computeRenderBounds(dim0, dim1, limits, limitScaling) {
    // Compute rendering bounds in data space
    return {
        xMin: limits.min[dim0] * limitScaling,
        xMax: limits.max[dim0] * limitScaling,
        yMin: limits.min[dim1] * limitScaling,
        yMax: limits.max[dim1] * limitScaling
    }
}

class CameraController extends THREE.Object3D {
    constructor(fixedCamera, dynamicCamera, dynamicCameraSpeed = 1) {
        super();
        this.activeCameraID = 0;  // 0 - fixedCamera, 1 - dynamicCamera
        this.fixedCamera = fixedCamera;
        this.dynamicCamera = dynamicCamera;
        this.dynamicCameraSpeed = dynamicCameraSpeed;
    }

    getActiveCamera() {
        if (this.activeCameraID === 0) {
            return this.fixedCamera;
        } else if (this.activeCameraID === 1) {
            return this.dynamicCamera;
        }
    }

    activateFixedCamera() {
        this.activeCameraID = 0;
        this.fixedCamera.visible = true;
        this.dynamicCamera.visible = false;
    }

    activateDynamicCamera() {
        this.activeCameraID = 1;
        this.fixedCamera.visible = false;
        this.dynamicCamera.visible = true;
    }

    cameraStep(dataset, deltaTime) {
        if (this.activeCameraID === 0) {
            // Do nothing
        } else if (this.activeCameraID === 1) {
            // move it, resize it
            let datasetCenter = dataset.pointsCenter();
            let stepSize = this.dynamicCameraSpeed * deltaTime / 1000;
            let targetX = datasetCenter.x;
            let targetY = datasetCenter.y;
            // this.dynamicCamera.position.x += stepSize * (datasetCenter.x - this.dynamicCamera.position.x);
            // this.dynamicCamera.position.y += stepSize * (datasetCenter.y - this.dynamicCamera.position.y);
            this.dynamicCamera.position.x = targetX;
            this.dynamicCamera.position.y = targetY;

            let datasetWidth = dataset.pointsWidth();
            let targetLeft = datasetCenter.x - datasetWidth / 2;
            let targetRight = datasetCenter.x + datasetWidth / 2;
            // this.dynamicCamera.left += stepSize * (this.dynamicCamera.left - targetLeft);
            // this.dynamicCamera.right += stepSize * (this.dynamicCamera.right - targetRight);
            // this.dynamicCamera.left = targetLeft;
            // this.dynamicCamera.right = targetRight;


            let datasetHeight = dataset.pointsHeight();
            let targetTop = datasetCenter.y - datasetHeight / 2;
            let targetBottom = datasetCenter.y + datasetHeight / 2;
            // this.dynamicCamera.top += stepSize * (this.dynamicCamera.top - targetTop);
            // this.dynamicCamera.bottom += stepSize * (this.dynamicCamera.bottom - targetBottom);
            this.dynamicCamera.top = targetTop;
            this.dynamicCamera.bottom = targetBottom;

            this.dynamicCamera.updateProjectionMatrix();
        }
    }

    setFixedCameraParameters(dim0, dim1, limits, limitScaling) {
        setCameraParameters(dim0, dim1, limits, limitScaling, this.fixedCamera);
    }
}

function valueLimits(dataStack) {
    // Compute minimum and maximum values for each dimension in the data stack across all time steps.
    let minima = [];
    let maxima = [];

    for (let step = 0; step < dataStack.length; step++) {
        for (let dim = 0; dim < dataStack[step].length; dim++) {
            if (minima.length <= dim) {
                minima.push(Infinity);
            }
            if (maxima.length <= dim) {
                maxima.push(-Infinity);
            }
            for (let particleId = 0; particleId < dataStack[step][dim].length; particleId++) {
                if (dataStack[step][dim][particleId] < minima[dim]) {
                    minima[dim] = dataStack[step][dim][particleId];
                }
                if (dataStack[step][dim][particleId] > maxima[dim]) {
                    maxima[dim] = dataStack[step][dim][particleId];
                }
            }
        }
    }

    return {
        'min': minima,
        'max': maxima
    };
}

// For DLA
class Dataset extends THREE.Object3D {
    constructor(
        data,
        r = 1,
        g = 1,
        b = 1,
        sprite = undefined,
        particleSize = 5,
        minima = undefined,  // TODO remove
        maxima = undefined,  // TODO remove
        z = 0
    ) {
        super();

        this.data = data;
        this.dim0 = null;
        this.dim1 = null;
        this.points = null;
        this.colors = null;
        this.z = z;

        this.minima = minima;  // list with nDimensions elements, minimum over all particles for each dimension
        this.maxima = maxima;  // list with nDimensions elements, maximum over all particles for each dimension

        this.geometry = new THREE.BufferGeometry();
        this.material = new THREE.PointsMaterial({
            size: particleSize,
            vertexColors: THREE.VertexColors,
            map: sprite,
            alphaTest: 0.5,
            transparent: true
        });
        this.points = new THREE.Points(this.geometry, this.material);

        this.add(this.points);  // Make this.points a child of this class

        this.setDimensions(0, 1)
        this.setColors(r, g, b);
    }

    pointsCenter() {
        return {
            x: this.minima[this.dim0] + this.pointsWidth() / 2,
            y: this.minima[this.dim1] + this.pointsHeight() / 2,
        }
    }

    pointsWidth() {
        return this.maxima[this.dim0] - this.minima[this.dim0];
    }

    pointsHeight() {
        return this.maxima[this.dim1] - this.minima[this.dim1];
    }

    setDimensions(dim0, dim1) {
        this.dim0 = dim0;
        this.dim1 = dim1;
        this.setPositions();
    }

    xs() {
        return this.data[this.dim0];
    }

    ys() {
        return this.data[this.dim1];
    }

    setColors(r, g, b) {
        this.colors = [];
        for (let i = 0; i < this.data[this.dim0].length; i++) {
            let color = new THREE.Color();
            color.setRGB(r, g, b);
            this.colors.push(color.r, color.g, color.b);
        }
        this.geometry.setAttribute('color', new THREE.Float32BufferAttribute(this.colors, 3));
    }

    setPositions() {
        let xs = this.xs();
        let ys = this.ys();
        let z = this.z;

        let positions = [];
        for (let i = 0; i < xs.length; i++) {
            positions.push(xs[i], ys[i], z);
        }
        this.geometry.setAttribute('position', new THREE.Float32BufferAttribute(positions, 3));
    }
}