# Copyright (c) 2019, Christoph Peters
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the Karlsruhe Institute of Technology nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import numpy as np
from matplotlib import pyplot
from scipy import optimize
from pickle import dump,load
from os import path

def GetWavelengthSampling():
    """Returns the wavelength sampling used by the XYZ curves."""
    return np.linspace(360.0,830.0,95);

def LoadXYZCurves():
    """Returns a quadruple (Lambda,X,Y,Z) of arrays where Lambda is the 
       wavelength in nanometers and X, Y and Z are the corresponding values of 
       the CIE 1931 X, Y, Z curves with arbitrary but consistent unit. Lambda 
       samples the interval from 360 nm to 830 nm (inclusive) at 5 nm steps."""
    # To make this code a bit more portable, we have the data in the code
    X=np.asarray([0.0001299,0.0002321,0.0004149,0.0007416,0.001368,0.002236,0.004243,0.00765,0.01431,0.02319,0.04351,0.07763,0.13438,0.21477,0.2839,0.3285,0.34828,0.34806,0.3362,0.3187,0.2908,0.2511,0.19536,0.1421,0.09564,0.05795001,0.03201,0.0147,0.0049,0.0024,0.0093,0.0291,0.06327,0.1096,0.1655,0.2257499,0.2904,0.3597,0.4334499,0.5120501,0.5945,0.6784,0.7621,0.8425,0.9163,0.9786,1.0263,1.0567,1.0622,1.0456,1.0026,0.9384,0.8544499,0.7514,0.6424,0.5419,0.4479,0.3608,0.2835,0.2187,0.1649,0.1212,0.0874,0.0636,0.04677,0.0329,0.0227,0.01584,0.01135916,0.008110916,0.005790346,0.004109457,0.002899327,0.00204919,0.001439971,0.0009999493,0.0006900786,0.0004760213,0.0003323011,0.0002348261,0.0001661505,0.000117413,0.00008307527,0.00005870652,0.00004150994,0.00002935326,0.00002067383,0.00001455977,0.00001025398,0.000007221456,0.000005085868,0.000003581652,0.000002522525,0.000001776509,0.000001251141]);
    Y=np.asarray([0.000003917,0.000006965,0.00001239,0.00002202,0.000039,0.000064,0.00012,0.000217,0.000396,0.00064,0.00121,0.00218,0.004,0.0073,0.0116,0.01684,0.023,0.0298,0.038,0.048,0.06,0.0739,0.09098,0.1126,0.13902,0.1693,0.20802,0.2586,0.323,0.4073,0.503,0.6082,0.71,0.7932,0.862,0.9148501,0.954,0.9803,0.9949501,1.0,0.995,0.9786,0.952,0.9154,0.87,0.8163,0.757,0.6949,0.631,0.5668,0.503,0.4412,0.381,0.321,0.265,0.217,0.175,0.1382,0.107,0.0816,0.061,0.04458,0.032,0.0232,0.017,0.01192,0.00821,0.005723,0.004102,0.002929,0.002091,0.001484,0.001047,0.00074,0.00052,0.0003611,0.0002492,0.0001719,0.00012,0.0000848,0.00006,0.0000424,0.00003,0.0000212,0.00001499,0.0000106,0.0000074657,0.0000052578,0.0000037029,0.0000026078,0.0000018366,0.0000012934,0.00000091093,0.00000064153,0.000000451810]);
    Z=np.asarray([0.0006061,0.001086,0.001946,0.003486,0.006450001,0.01054999,0.02005001,0.03621,0.06785001,0.1102,0.2074,0.3713,0.6456,1.0390501,1.3856,1.62296,1.74706,1.7826,1.77211,1.7441,1.6692,1.5281,1.28764,1.0419,0.8129501,0.6162,0.46518,0.3533,0.272,0.2123,0.1582,0.1117,0.07824999,0.05725001,0.04216,0.02984,0.0203,0.0134,0.008749999,0.005749999,0.0039,0.002749999,0.0021,0.0018,0.001650001,0.0014,0.0011,0.001,0.0008,0.0006,0.00034,0.00024,0.00019,0.0001,0.00004999999,0.00003,0.00002,0.00001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]);
    return (GetWavelengthSampling(),X,Y,Z);


def ApplyWarpDensity(Wavelength,WarpDensity,MirrorSignal):
    """This function maps the given wavelengths to phases using the specified 
       warp."""
    CDF=WarpDensity.cumsum()-WarpDensity[0];
    CDF/=CDF[CDF.size-1];
    return np.interp(Wavelength,GetWavelengthSampling(),CDF)*((1.0 if MirrorSignal else 2.0)*np.pi)-np.pi;


def GetWarpedFourierBasis(WarpDensity,MirrorSignal,nDimension=3):
    """Models the effect of integrating over a Fourier basis that is warped by 
       a piecewise linear function with the given derivative. The scaling is 
       arbitrary."""
    Phase=ApplyWarpDensity(GetWavelengthSampling(),WarpDensity,MirrorSignal);
    WarpDensity[0]=0.0;
    Derivative=np.concatenate([[0.0],WarpDensity[1:]])+np.concatenate([WarpDensity[1:],[0.0]]);
    # Warp Fourier
    if(MirrorSignal):
        Fourier=np.vstack([np.cos(j*Phase)*Derivative for j in range(nDimension)]).T;
    else:
        Fourier=np.vstack([WarpDensity]+[np.cos(j*Phase)*Derivative for j in range(1,(nDimension+1)//2)]+[np.sin(j*Phase)*WarpDensity for j in range(1,(nDimension+1)//2)]).T;
    return Fourier;


def GetOrthogonalProjector(Matrix):
    """Constructs and returns an orthogonal projector onto the column space of 
       the given matrix."""
    MatrixOrthonormal=np.linalg.qr(Matrix)[0];
    return MatrixOrthonormal@MatrixOrthonormal.T;


def GetWarpFilePath(UseMirroring,nSample,ColorSpaceName):
    """Returns the path to the file where the warp for the specified parameter 
       set will be stored."""
    WarpOutputDirectory="./";
    return path.join(WarpOutputDirectory,"%sWarp%s%d.pickle"%(ColorSpaceName,"Even" if UseMirroring else "",nSample));


def OptimizeFunctionSpaceDistance(UseMirroring,X=None,Y=None,Z=None,ColorSpaceName="XYZ"):
    """Optimizes a warp to minimize the distance between the Fourier basis and 
       the color space spanned by the given three response curves. They should 
       be sampled at the wavelengths given by GetWavelengthSampling(). If they 
       are None, CIE XYZ is used. The result is dumped to a file for subsequent 
       use, which includes the given string in the file name.
      \warning This function uses non-convex optimization. It has worked as is 
               for XYZ but if you want to optimize for other response curves, 
               you might have to fiddle with the optimization technique or the 
               parameters. In any case, you should inspect the result 
               carefully."""
    # Sanitize the inputs
    if(X is None):
        X,Y,Z=LoadXYZCurves()[1:];
    else:
        X=np.asarray(X);
        Y=np.asarray(Y);
        Z=np.asarray(Z);
    nSample=X.size;
    # Construct an orthonormal projector onto XYZ
    XYZ=np.vstack([X,Y,Z]).T;
    ProjectorXYZ=GetOrthogonalProjector(XYZ);
    def GetFunctionSpaceDistance(WarpDensity):
        Fourier=GetWarpedFourierBasis(WarpDensity,UseMirroring);
        # Construct an orthonormal projector for Fourier
        FourierOrthonormal=np.linalg.qr(Fourier)[0];
        ProjectorFourier=FourierOrthonormal@FourierOrthonormal.T;
        # Evaluate the error
        return np.linalg.norm(ProjectorXYZ-ProjectorFourier,2);
    # A bunch of solver options and solvers are commented out below. If you 
    # want to experiment, these are good starting points. basinhopping is a 
    # global optimizer using a given local optimizer internally. It is much 
    # slower but also much more powerful.
    #optimize.show_options("minimize","L-BFGS-B");
    Initialization=XYZ.sum(1);
    FunctionErrorTolerance=1.0e-13;
    SolverDict=dict(method="L-BFGS-B",options=dict(maxfun=1000000,ftol=FunctionErrorTolerance));
    #SolverDict=dict(method="SLSQP",options=dict(maxiter=1000000,ftol=FunctionErrorTolerance));
    SolverDict["bounds"]=[(0.0,None)]*nSample;
    Result=optimize.minimize(GetFunctionSpaceDistance,Initialization,**SolverDict);
    #Result=optimize.basinhopping(GetFunctionSpaceDistance,Initialization,minimizer_kwargs=SolverDict,disp=True);
    WarpDensity=Result.x;
    dump(WarpDensity,open(GetWarpFilePath(UseMirroring,nSample,ColorSpaceName)+("" if Result.success else ".failed"),"wb"));
    print(Result.message);


def PlotFunctionSpaceDistance(MirrorSignal,X=None,Y=None,Z=None,ColorSpaceName="XYZ"):
    """Plots the projection of the given basis into the space spanned by a 
       Fourier basis with a warp computed and stored previously by 
       OptimizeFunctionSpaceDistance(). Arguments are the same. Additionally it 
       visualizes the used error metric of the optimization."""
    # Sanitize the inputs
    Lambda=GetWavelengthSampling();
    if(X is None):
        X,Y,Z=LoadXYZCurves()[1:];
    else:
        X=np.asarray(X);
        Y=np.asarray(Y);
        Z=np.asarray(Z);
    nSample=Lambda.size;
    if(X.size!=nSample or Y.size!=nSample or Z.size!=nSample):
        print("X, Y and Z must be sampled at the wavelengths returned by GetWavelengthSampling().");
        return;
    # Construct orthogonal projectors for the original space and the warped 
    # Fourier space
    XYZ=np.vstack([X,Y,Z]).T;
    ProjectorXYZ=GetOrthogonalProjector(XYZ);
    WarpFilePath=GetWarpFilePath(UseMirroring,nSample,ColorSpaceName);
    WarpDensity=load(open(WarpFilePath,"rb"));
    Fourier=GetWarpedFourierBasis(WarpDensity,MirrorSignal,3);
    FourierOrthonormal=np.linalg.qr(Fourier)[0];
    ProjectorFourier=FourierOrthonormal@FourierOrthonormal.T;
    # Plot the original basis and its projection into the warped Fourier basis
    for i in range(3):
        pyplot.subplot(2,2,i+1,title="Fit to %s"%("XYZ"[i]),xlim=(Lambda.min(),Lambda.max()),xlabel="Wavelength / nm");
        pyplot.plot(Lambda,XYZ[:,i]);
        pyplot.plot(Lambda,ProjectorFourier@XYZ[:,i]);
    # Analyze the distance. Since our distance measure is based on the matrix 
    # 2-norm (also known as spectral norm), the error can be analyzed by 
    # considering a singular vector for the largest singular value.
    pyplot.subplot(2,2,4,title="Error",xlim=(Lambda.min(),Lambda.max()),xlabel="Wavelength / nm");
    ProjectorDifference=ProjectorXYZ-ProjectorFourier;
    print("Error: ",np.linalg.norm(ProjectorDifference,2));
    ErrorVector=np.linalg.svd(ProjectorDifference)[2][0,:];
    pyplot.plot(Lambda,ErrorVector,label="Error vector",linewidth=3.0);
    pyplot.plot(Lambda,ProjectorXYZ@ErrorVector,label="XYZ projection");
    pyplot.plot(Lambda,ProjectorFourier@ErrorVector,label="Fourier projection");
    pyplot.plot(Lambda,WarpDensity/np.sum(WarpDensity),label="Warp density");
    pyplot.legend(loc="lower right",frameon=False);
    pyplot.tight_layout();


def ApplyWarp(MirrorSignal,Wavelength,ColorSpaceName="XYZ"):
    """This function applies the warp produced by minimizing the distance between 
       function spaces. It maps the given wavelengths to phases."""
    WarpDensity=load(open(GetWarpFilePath(MirrorSignal,95,ColorSpaceName),"rb"));
    return ApplyWarpDensity(Wavelength,WarpDensity,MirrorSignal);

def ApplyXYZWarpEven(Wavelength):
    """Short-hand for ApplyWarp(True,Wavelength,"XYZ")."""
    return ApplyWarp(True,Wavelength);

def ApplyXYZWarpOdd(Wavelength):
    """Short-hand for ApplyWarp(True,Wavelength,"XYZ")."""
    return ApplyWarp(False,Wavelength);


if(__name__=="__main__"):
    # Here you can specify what situation you want to optimize the warp for. 
    # Note that you do not need to rerun the optimization if your color space 
    # is defined in terms of XYZ. In that case a simple linear transform will 
    # suffice. You only need to redo it, if you have, for example, care about 
    # response curves measured for a particular digital camera.
    UseMirroring=False;
    X,Y,Z=LoadXYZCurves()[1:];
    ColorSpaceName="XYZ";
    # Perform the optimization
    OptimizeFunctionSpaceDistance(UseMirroring,X,Y,Z,ColorSpaceName);
    # Visualize the result
    PlotFunctionSpaceDistance(UseMirroring,X,Y,Z,ColorSpaceName);
    # Print the phases for the usual CIE wavelength sampling. This may be copy-
    # pasted into some program as lookup table.
    print(", ".join(["%.9ef"%Phase for Phase in ApplyWarp(UseMirroring,GetWavelengthSampling())]));
    pyplot.show();
    print("Done.");
