package nu.sutic.ncp;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.File;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;
import java.util.ArrayList;


/**
 * A Nikon Custom Picture Control.
 */
public class NikonCustomPictureControl {
    // Record 1 Parameters
    public String version = "0100";
    public String name = "CUSTOM";
    public BaseProfile baseProfile = BaseProfile.NEUTRAL; // Default NEUTRAL
    public int modified = 2;          // Default Modified

    public int unknown1 = 0;

    public int sharpening = 0;
    public int contrast = 0;
    public int brightness = 0;
    public int saturation = 0;
    public int hue = 0;
    public MonoFilter filter = MonoFilter.NOT_APPLICABLE;
    public MonoToning toning = MonoToning.NOT_APPLICABLE;
    public int toningStrength = 0;

    // Record 2 Parameters (Curve Data)
    public short unknown2 = 0;

    /**
     * Black point used for midtone gamma calculations.
     */
    public int blackPoint = 0;

    /**
     * White point used for midtone gamma calculations.
     */
    public int whitePoint = 255;

    /**
     * Integer portion of the halftone gamma value. Actual gamma value is
     *
     *     halftoneGammaInteger + halftoneGammaFractionPercent * 0.01
     *
     */
    public int halftoneGammaInteger = 1;

    /**
     * Fractional part of the halftone gamma value in units
     * of 0.01. Actual gamma value is
     *
     *     halftoneGammaInteger + halftoneGammaFractionPercent * 0.01
     *
     */
    public int halftoneGammaFractionPercent = 0;

    /**
     * Sets the global halftone (midtone gamma) value for the profile.
     * @param value Typically a value between 0.50 and 2.00 (Default is 1.00)
     */
    public void setHalftoneGamma(double value) {
        // Example: 1.45 -> halftoneGammaInteger = 1, halftoneGammaFractionPercent = 45
        this.halftoneGammaInteger = (int) Math.floor(value);
        this.halftoneGammaFractionPercent = (int) Math.round((value - this.halftoneGammaInteger) * 100.0);
    }

    /**
     * Gets the readable halftone value from the profile state.
     */
    public double getHalftoneGamma() {
        return this.halftoneGammaInteger + (0.01 * this.halftoneGammaFractionPercent);
    }

    /**
     * The lowest output level. the [lut] is clamped between this and [outputMax]
     * except for the last entry.
     */
    public int outputMin = 0;

    /**
     * The highest output level. the [lut] is clamped between this and [outputMax]
     * except for the last entry.
     */
    public int outputMax = 255;

    /**
     * Number of spline curve points.
     */
    public int numPoints = 0;

    /**
     * [numPoints] user-interface points. Maximum 19 points represented as two bytes. These are used
     * for the UI and to compute the [lut] field, which is the one used
     * by the camera to transform the image.
     */
    public List<CurvePoint> curvePoints = new ArrayList<>();

    /**
     * Following the [numPoints] curvePoints are padding bytes so that the
     * curvePoints plus the padding are a total of 57 bytes. That is, there
     * are `57 - 2 * numPoints` bytes in the array.
     */
    public byte[] curvePointsPadding;

    /**
     * Actual 1D-LUT used for image transformation. This array
     * stores the lut values for inputs 1 - 256. Input value zero is not
     * stored in the array.
     */
    public int[] lut = new int[256];

    public NikonCustomPictureControl() {}

    /**
     * Recomputes the [lut] field.
     *
     * The [lut] field has 256 entries and they correspond to input brightnesses of
     * one to 256, inclusive.
     *
     * The easiest way to look at what is happening is to see the
     * LUT as two *parallel* modifiers to the no-op NCP curve, which would be
     * a linear slope from zero to max output as the input goes from
     * zero to max input.
     * To clarify, we're not doing $`y = Spline(Gamma(x))`$ or $`y = Gamma(Spline(x))`$,
     * but instead a sort of $`y = Spline(x) + Gamma(x)`$.
     *
     * If we look at the two functions as corrections to the linear no-op curve we have:
     *
     *  * A midtone gamma correction - the difference between the gamma-corrected
     *    brightness curve and the linear one: $`Gamma(x) - NoOp(x)`$
     *  * A spline correction - the difference between the spline and the
     *    linear curve: $`Spline(x) - NoOp(x)`$
     *  * The linear curve $`NoOp(x)`$ which is just $`x`$.
     *
     * The 1D-LUT is then:
     *
     * $$
     * \text{output} = (\text{linear curve}) + \\
     *                 (\text{midtone gamma correction}) + \\
     *                 (\text{spline correction})
     * $$
     *
     * or, to spell out the functions as is done in the code below:
     *
     * $$
     * y = NoOp(x) + (Gamma(x) - NoOp(x)) + (Spline(x) - NoOp(x))\\
     * y = x + (Gamma(x) - x) + (Spline(x) - x)
     * $$
     *
     * This simplifies to:
     *
     * $$
     * y = Gamma(x) + Spline(x) - x
     * $$
     *
     * but for clarity the code does it the roundabout way.
     */
    public void recomputeLut() {
        // Set up the spline.
        //
        // For unknown reasons, the X axis must be stretched to
        // 256.0 - possibly because we only have an 8-bit value
        // to specify the white point and we need to extend it
        // to 255.9999 to accomodate the higher-bit input values
        // of the sensor data.
        List<NaturalSpline.Point> handles = new ArrayList<>();
        for (CurvePoint p : curvePoints) {
            handles.add(new NaturalSpline.Point(
                p.input() * 256.0 / 255.0,
                p.output()
            ));
        }
        NaturalSpline spline = new NaturalSpline(handles);
        var firstHandle = spline.first();
        var lastHandle = spline.last();

        // Input starts at 1, because we do not store the
        // first entry in the LUT array - the X value for the
        // first array entry is 1.0, not 0.0.
        for (int inputX = 1; inputX <= 255; inputX++) {
            double sampleX = inputX;

            double output = 0.0;
            if (sampleX < firstHandle.x()) {
                // If we're before the spline, assume same
                // y as the first handle.
                output = firstHandle.y();
            } else if (sampleX >= lastHandle.x()) {
                // Same if we're after the spline, assume same
                // y as the last handle.
                output = lastHandle.y();
            } else {
                // We're inside the spline region,
                // evaluate it
                double splineY = spline.evaluate(sampleX);

                double normalizedX = normalizedInput(sampleX);
                double pureGamma = Math.pow(normalizedX, 1.0 / getHalftoneGamma()) * getOutputRange();
                double linearSlope = normalizedX * getOutputRange();

                /// Output is the linear slope (just $`x`$, no corrections), plus
                // the "gamma correction", which is $`gamma(x) - x`$, plus
                // the "spline correction" to the linear slope, which is $`spline(x) - x`$.
                // We spell it out for clarity even if it's suboptimal.
                output = linearSlope + (pureGamma - linearSlope) + (splineY - linearSlope);
            }

            // Clamp output to the min and max output levels
            if (output < outputMin) {
                output = outputMin;
            } else if (output > outputMax) {
                output = outputMax;
            }

            // Map value using Nikon's mandatory x128 LUT scalar representation
            // that maps the 8-bit range to 15 bits.
            lut[inputX - 1] = (int) Math.round(output * 32767.0 / 255.0);
        }

        // The last entry appears to be fixed to this.
        lut[255] = (int) Math.round(lastHandle.y() * 32767.0 / 255.0);
    }

    /**
     * The range from outputMax to outputMin.
     */
    public int getOutputRange() {
        return outputMax - outputMin;
    }

    /**
     * The range from blackPoint to whitePoint.
     */
    public int getInputRange() {
        return whitePoint - blackPoint;
    }

    /**
     * Remaps a [0 ... 255] input value to [0 ... 1] with
     * 0.0 being [blackPoint] and 1.0 being [whitePoint].
     */
    public double normalizedInput(double x) {
        return (x - blackPoint) / getInputRange();
    }

    /**
     * Read an NCP profile from a file.
     */
    public static NikonCustomPictureControl read(File filePath) throws IOException {
        try (InputStream is = new FileInputStream(filePath)) {
            return read(is);
        }
    }

    /**
     * Read from an input stream. The stream is expected to be positioned
     * at the start of the profile.
     */
    public static NikonCustomPictureControl read(InputStream input) throws IOException {
        NikonCustomPictureControl profile = new NikonCustomPictureControl();

        DataInputStream dis = new DataInputStream(input);
        byte[] signature = new byte[4];
        dis.readFully(signature);
        if (!"NCP\0".equals(new String(signature, StandardCharsets.US_ASCII))) {
            throw new IOException("Invalid file signature. Not a valid Nikon NCP profile.");
        }

        while (true) {
            int recordID;
            try {
                recordID = dis.readInt();
            } catch (IOException e) {
                break; // Natural EOF
            }

            if (recordID == 0) break;
            int recordSize = dis.readInt();

            byte[] recordData = new byte[recordSize];
            dis.readFully(recordData);
            ByteBuffer buffer = ByteBuffer.wrap(recordData).order(ByteOrder.BIG_ENDIAN);

            if (recordID == 1) {
                byte[] verBytes = new byte[4];
                buffer.get(verBytes);
                profile.version = new String(verBytes, StandardCharsets.US_ASCII);

                byte[] nameBytes = new byte[20];
                buffer.get(nameBytes);
                profile.name = new String(nameBytes, StandardCharsets.US_ASCII).split("\0")[0].trim();

                buffer.mark();
                profile.baseProfile = BaseProfile.of(buffer.getShort() & 0xFFFF);
                if (profile.baseProfile == null) {
                    // Could be a real old profile with corrupt encoding.
                    // Try parsing the base profile as a single byte
                    // and use ofNOP.
                    buffer.reset();
                    profile.baseProfile = BaseProfile.ofNOP(buffer.get() & 0xff);
                }
                if (profile.baseProfile == null) {
                    throw new IllegalArgumentException("Unknown base profile");
                }
                profile.modified = buffer.get() & 0xFF;
                profile.unknown1 = buffer.get() & 0xFF;

                profile.sharpening = (buffer.get() & 0xFF) - 0x80;
                profile.contrast = (buffer.get() & 0xFF) - 0x80;
                profile.brightness = (buffer.get() & 0xFF) - 0x80;
                profile.saturation = (buffer.get() & 0xFF) - 0x80;
                profile.hue = (buffer.get() & 0xFF) - 0x80;

                int filterIndex = (buffer.get() & 0xFF) - 0x80;
                profile.filter = MonoFilter.of(filterIndex);
                profile.toning = MonoToning.of((buffer.get() & 0xFF) - 0x80);
                profile.toningStrength = (buffer.get() & 0xFF) - 0x80;
            } else if (recordID == 2) {
                profile.unknown2 = buffer.getShort();
                profile.blackPoint = buffer.get() & 0xFF;
                profile.whitePoint = buffer.get() & 0xFF;

                profile.outputMin = buffer.get() & 0xFF;
                profile.outputMax = buffer.get() & 0xFF;

                profile.halftoneGammaInteger = buffer.get() & 0xFF;
                profile.halftoneGammaFractionPercent = buffer.get() & 0xFF;

                profile.numPoints = buffer.get() & 0xFF;

                // Now follows 57 bytes which contain
                // numPoints two-byte curve points.
                for (int i = 0; i < profile.numPoints; i++) {
                    int inputValue = buffer.get() & 0xFF;
                    int outputValue = buffer.get() & 0xFF;
                    profile.curvePoints.add(new CurvePoint(inputValue, outputValue));
                }
                // The remaining bytes of the 57 bytes are skipped.
                for (int skip = 0; skip < (57 - profile.numPoints * 2); ++skip) {
                    buffer.get();
                }

                // Then comes 256 LUT entries.
                for (int i = 0; i < 256; i++) {
                    profile.lut[i] = buffer.getShort() & 0xFFFF;
                }
            }
        }

        // Some curves have a black point of zero, even though the spline
        // drops to zero a bit above that. The result is that midtone gamma
        // is applied to the region below, resulting in reconstruction
        // errors.
        //
        // We detect the blackPoint == 0 and reset it to the point where the curve
        // drops to zero.
        if (profile.blackPoint == 0 && profile.curvePoints.get(0).output() == 0) {
            profile.blackPoint = profile.curvePoints.get(0).input();
        }

        // Same for the white point, except now we look at the last curve point
        // and if it hits 255.
        if (profile.whitePoint == 255 && profile.curvePoints.get(profile.curvePoints.size() - 1).output() == 255) {
            profile.whitePoint = profile.curvePoints.get(profile.curvePoints.size() - 1).input();
        }
        return profile;
    }

    public void write(File file) throws IOException {
        try (OutputStream os = new FileOutputStream(file)) {
            write(os);
        }
    }

    public void write(OutputStream out) throws IOException {
        DataOutputStream dos = new DataOutputStream(out);
        // Write Magic Header Signature
        dos.write("NCP\0".getBytes(StandardCharsets.US_ASCII));

        // ---- PACK RECORD 1 ----
        dos.writeInt(1); // Record ID
        dos.writeInt(36); // Fixed size allocation for record 1 fields

        byte[] verBytes = Arrays.copyOf(version.getBytes(StandardCharsets.US_ASCII), 4);
        dos.write(verBytes);

        byte[] nameBytes = Arrays.copyOf(name.getBytes(StandardCharsets.US_ASCII), 20);
        dos.write(nameBytes);

        dos.writeShort(baseProfile.value);
        dos.writeByte(modified);
        dos.writeByte(unknown1);

        dos.writeByte((sharpening + 0x80) & 0xFF);
        dos.writeByte((contrast + 0x80) & 0xFF);
        dos.writeByte((brightness + 0x80) & 0xFF);
        dos.writeByte((saturation + 0x80) & 0xFF);
        dos.writeByte((hue + 0x80) & 0xFF);
        dos.writeByte((filter.value + 0x80) & 0xFF);
        dos.writeByte((toning.value + 0x80) & 0xFF);
        dos.writeByte((toningStrength + 0x80) & 0xFF);

        // ---- PACK RECORD 2 ----
        dos.writeInt(2); // Record ID
        dos.writeInt(578); // Struct block allocation size: 2+1+1+1+1+1+1+1+57+(256*2)

        dos.writeShort(unknown2);

        dos.writeByte(blackPoint);
        dos.writeByte(whitePoint);

        dos.writeByte(outputMin);
        dos.writeByte(outputMax);

        dos.writeByte(halftoneGammaInteger);
        dos.writeByte(halftoneGammaFractionPercent);

        dos.writeByte(numPoints);
        for (CurvePoint p : curvePoints) {
            dos.writeByte(p.input());
            dos.writeByte(p.output());
        }

        for (int skip = 0; skip < (57 - numPoints * 2); ++skip) {
            dos.writeByte(0);
        }

        for (int i = 0; i < 256; i++) {
            dos.writeShort(lut[i]);
        }

        // ---- EOF RECORD ----
        dos.writeInt(0);
    }
}