/*

  xmunipack - color space

  Copyright © 2009-2014, 2019 F.Hroch (hroch@physics.muni.cz)

  This file is part of Munipack.

  Munipack is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  Munipack is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Munipack.  If not, see <http://www.gnu.org/licenses/>.


  Color images:

  rawtran -o IMG_5807.fits -X "-q 3 -n 500" IMG_5807.CR2

 * important: -q selects adequate interpolation method,
              -n 500 selects threshold for wavelets


*/


#include "fits.h"
#include <wx/wx.h>
#include <wx/wfstream.h>
#include <wx/txtstrm.h>
#include <wx/tokenzr.h>

/* CIE Lab, D65 white point */
const float FitsColor::Xn =  95.047;
const float FitsColor::Yn = 100.000;
const float FitsColor::Zn = 108.883;


// ---- reference counting data base
class FitsColorData : public wxObjectRefData
{
public:
  FitsColorData();
  FitsColorData(const FitsColorData&);
  FitsColorData& operator = (const FitsColorData&);
  virtual ~FitsColorData();

  int ncolors, nbands;
  float *trafo,*level,*weight;
};



FitsColorData::FitsColorData()
{
  ncolors = 3;
  nbands = 3;
  trafo = new float[ncolors*nbands];
  level = new float[ncolors];
  weight = new float[ncolors];
  for(int i = 0; i < ncolors; i++) {
    level[i] = 0.0;
    weight[i] = 1.0;
    for(int j = 0; j < nbands; j++)
      trafo[j*ncolors + i] = i == j ? 1.0 : 0.0;
  }
}


FitsColorData::FitsColorData(const FitsColorData& copy)
{
  wxFAIL_MSG("FitsColorData WE ARE REALY NEED COPY CONSTRUCTOR");
}

FitsColorData& FitsColorData::operator = (const FitsColorData& other)
{
  wxFAIL_MSG("FitsColorData: WE ARE REALLY NEED ASSIGNMENT CONSTRUCTOR");
  return *this;
}


FitsColorData::~FitsColorData()
{
  delete[] trafo;
  delete[] level;
  delete[] weight;
}


// --- FitsColor


FitsColor::FitsColor():
  ispace(COLORSPACE_XYZ),saturation(1.0),hue(0.0),
  nitevision(false),nitethresh(0.0),nitewidth(10.0)
{
  UnRef();
  SetRefData(new FitsColorData);
}

FitsColor::FitsColor(const wxString& cdatafile, const FitsArray& array):
  ispace(COLORSPACE_XYZ),saturation(1.0),hue(0.0),
  nitevision(false),nitethresh(0.0),nitewidth(10.0)
{
  UnRef();
  SetRefData(new FitsColorData);

  Init(cdatafile,array);
}

//FitsColor::~FitsColor() { UnRef(); }

void FitsColor::Init(const wxString& cdatafile, const FitsArray& array)
{
  wxString cs = array.GetKey(FITS_KEY_CSPACE);
  if( cs.IsEmpty() )
    cs = "RGB";
  // RGB is supposed due to compatibility with other software


  SetTrans(cs,cdatafile);

  for(int i = 0; i < array.Naxes(2); i++) {
    FitsArrayStat s(array.Plane(i));
    SetLevel(i,s.Med());
  }
}

float FitsColor::GetWeight(int n) const
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->nbands);
  return data->weight[n];
}

float FitsColor::GetLevel(int n) const
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->nbands);
  return data->level[n];
}


void FitsColor::SetWeight(int n, float w)
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->nbands);
  data->weight[n] = w;
}

void FitsColor::SetLevel(int n, float x)
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->nbands);
  data->level[n] = x;
}

void FitsColor::SetTrans(int n, int m, float x)
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->ncolors && 0 <= m && m < data->nbands);
  data->trafo[m*data->nbands + n] = x;
}


void FitsColor::SetTrans(int n, int m)
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data);
  delete[] data->weight;
  delete[] data->level;
  delete[] data->trafo;

  data->ncolors = n;
  data->nbands = m;
  data->weight = new float[n];
  data->level = new float[n];
  data->trafo = new float[n*m];
  for(int i = 0; i < n; i++) {
    data->weight[i] = 1.0;
    data->level[i] = 0.0;
    for(int j = 0; j < m; j++)
      data->trafo[j*n + i] = i == j ? 1.0 : 0.0;
  }
}


void FitsColor::SetTrans(const wxString& cs)
{
  cspace = cs;

  if( cspace.Find("XYZ") != wxNOT_FOUND ) {
    size_t n = 3;
    SetTrans(n,n);
    for(size_t i = 0; i < n; i++)
      for(size_t j = 0; j < n; j++)
	SetTrans(i,j,0.0);

    for(size_t i = 0; i < n; i++) {
      SetTrans(i,i,1.0);
      SetLevel(i,0.0);
      SetWeight(i,1.0);
    }
  }
}

void FitsColor::SetTrans(const wxString& cs, const wxString& filename)
{
  cspace = cs;

  wxFileInputStream input(filename);
  wxTextInputStream text(input," ,\t");

  while(input.IsOk() && ! input.Eof()) {
    wxString line,ilabel,olabel;
    int n,m;

    line = text.ReadLine();
    line.Trim();
    if( line.IsEmpty() ) continue;

    wxArrayString a;
    wxStringTokenizer t(line,"'");
    int i = 0;
    while ( t.HasMoreTokens() ) {
      wxString x = t.GetNextToken();
      x.Trim();
      if( ! x.IsEmpty() ) {
	a.Add(x);
      }
      i++;
    }

    if( a.GetCount() == 2 ) {
      ilabel = a[0];
      olabel = a[1];
    }

    if( input.Eof() ) break;
    n = text.Read32();
    if( input.Eof() ) break;
    m = text.Read32();
    if( input.Eof() ) break;

    float *cmatrix = new float[n*m];
    for(int i = 0; i < n*m; i++) {
      if( input.Eof() ) break;
      cmatrix[i]= text.ReadDouble();
    }

    //   wxLogDebug(ilabel + " > " +olabel+ " , " +cspace+ " , " + Type_str(COLORSPACE_XYZ));

    if( ilabel == cspace && olabel == Type_str(COLORSPACE_XYZ) ) {
      FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
      wxASSERT(data);
      delete[] data->weight;
      delete[] data->level;
      delete[] data->trafo;

      data->ncolors = n;
      data->nbands = m;
      data->weight = new float[n];
      data->level = new float[n];
      data->trafo = cmatrix;
      for(int i = 0; i < n; i++) {
	data->weight[i] = 1.0;
	data->level[i] = 0.0;
      }

      return;
    }
    delete[] cmatrix;
  }


  // proper colorspace data not found
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data);
  delete[] data->weight;
  delete[] data->level;
  delete[] data->trafo;

  const int n = 3;
  const int m = 3;
  data->ncolors = n;
  data->nbands = m;
  data->weight = new float[n];
  data->level = new float[n];
  data->trafo = new float[n*m];
  for(int i = 0; i < n; i++) {
    data->weight[i] = 1.0;
    data->level[i] = 0.0;
    for(int j = 0; j < m; j++)
      data->trafo[j*n + i] = i == j ? 1.0 : 0.0;
  }
}


int FitsColor::GetColors() const
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data);
  return data->ncolors;
}

int FitsColor::GetBands() const
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data);
  return data->nbands;
}

float FitsColor::GetTrans(int n, int m) const
{
  FitsColorData *data = static_cast<FitsColorData *>(GetRefData());
  wxASSERT(data && 0 <= n && n < data->ncolors && 0 <= m && m < data->nbands);
  return *(data->trafo+m*data->nbands + n);
}

wxString FitsColor::GetColorspace() const
{
  return cspace;
}


void FitsColor::SetSaturation(float x)
{
  saturation = x;
}

void FitsColor::SetHue(float x)
{
  hue = x/57.29577951;
}

void FitsColor::SetNiteThresh(float x)
{
  nitethresh = x;
}

void FitsColor::SetNiteWidth(float x)
{
  nitewidth = x;
}

void FitsColor::SetNiteVision(bool t)
{
  nitevision = t;
}

void FitsColor::Reset()
{
  ispace = COLORSPACE_XYZ;
  saturation = 1;
  hue = 0;
  nitethresh = 0;
  nitewidth = 10;
  nitevision = false;
}

float FitsColor::NiteProfile(float x) const
{
  return 1.0/(1.0 + expf(-2.5*(x - nitethresh)/nitewidth));
}

wxString FitsColor::Type_str(int n)
{
  switch(n){
  case COLORSPACE_XYZ:      return "XYZ";
  default:        return wxEmptyString;
  }
}


wxArrayString FitsColor::Type_str()
{
  wxArrayString a;

  for(int i = COLORSPACE_XYZ+1; i < COLORSPACE_LAST; i++)
    a.Add(Type_str(i));

  return a;
}


float FitsColor::Scotopic(float X, float Y, float Z)
{
  return 0.36169*Z + 1.18214*Y - 0.80498*X;
}


float FitsColor::InvGamma(float r)
{
  if( r < 0.03928 )
    return r/12.92;
  else
    return powf((r + 0.055)/1.055,2.4);
}


void FitsColor::Instr_XYZ(long npix, size_t nband, const float **d,
			  float *Z, float *Y, float *X)
{
  wxASSERT(nband == (size_t) GetBands());

  int ncolors = GetColors();
  int nbands = GetBands();

  long nbytes = npix*sizeof(float);
  memcpy(X,d[0],nbytes);
  memcpy(Y,d[1],nbytes);
  memcpy(Z,d[2],nbytes);
  return;

  if( GetColorspace().Find("XYZ") != wxNOT_FOUND ) {
    wxASSERT(ncolors == nbands);

    long nbytes = npix*sizeof(float);
    memcpy(X,d[0],nbytes);
    memcpy(Y,d[1],nbytes);
    memcpy(Z,d[2],nbytes);
  }
  else {

    // allocate all on heap?

    float cb[ncolors][nbands];
    for(int i = 0; i < ncolors; i++)
      for(int j = 0; j < nbands; j++) {
	cb[i][j] = GetTrans(i,j);
      }

    float weight[nbands], level[nbands];
    for(int j = 0; j < nbands; j++) {
      weight[j] = GetWeight(j);
      level[j] = GetLevel(j);
    }

    float b[nbands];
    float c[ncolors];
    for(long i = 0; i < npix; i++) {

      for(int j = 0; j < nbands; j++) {
	b[j] = weight[j]*(d[j][i] - level[j]);
	if( b[j] < 0.0 ) b[j] = 0.0;
      }

      for(int l = 0; l < ncolors; l++) {
	float s = 0.0;
	for(int j = 0; j < nbands; j++)
	  s = s + cb[l][j]*b[j];
	c[l] = s;
      }

      X[i] = c[0];
      Y[i] = c[1];
      Z[i] = c[2];
    }
  }

}


// https://en.wikipedia.org/wiki/CIELAB_color_space
void FitsColor::XYZ_Lab(long npix, float *X, float *Y, float *Z,
			float *L, float *a, float *b)
{
  float *fX = new float[npix];
  float *fY = new float[npix];
  float *fZ = new float[npix];

  Lab_fun(npix,X,Xn,fX);
  Lab_fun(npix,Y,Yn,fY);
  Lab_fun(npix,Z,Zn,fZ);

  for(long i = 0; i < npix; i++) {
    L[i] = 116.0f*fY[i] - 16.0f;
    a[i] = 500.0f*(fX[i] - fY[i]);
    b[i] = 200.0f*(fY[i] - fZ[i]);
  }

  delete[] fX;
  delete[] fY;
  delete[] fZ;
}

void FitsColor::Lab_fun(long npix, float *I, float In, float *f)
{
  float d = 6.0 / 29.0;
  float d2 = d*d;
  float d3 = d2*d;
  float a = 3.0*d2;
  float b = 4.0 / 29.0;

  for(long i = 0; i < npix; i++) {
    float r = I[i] / In;
    if( r > d3 )
      f[i] = cbrtf(r);
    else
      f[i] = r / a + b;
  }
}

void FitsColor::Lab_XYZ(long npix, float *L, float *a, float *b,
			float *X, float *Y, float *Z)
{
  float *Lx = new float[npix];
  float *ax = new float[npix];
  float *bx = new float[npix];

  for(long i = 0; i < npix; i++) {
    Lx[i] = (L[i] + 16.0f) / 116.0f;
    ax[i] = Lx[i] + a[i] / 500.0f;
    bx[i] = Lx[i] - b[i] / 200.0f;
  }

  Lab_invfun(npix,ax,Xn,X);
  Lab_invfun(npix,Lx,Yn,Y);
  Lab_invfun(npix,bx,Zn,Z);

  delete[] Lx;
  delete[] ax;
  delete[] bx;
}

void FitsColor::Lab_invfun(long npix, float *f, float In, float *I)
{
  float d = 6.0 / 29.0;
  float a = 3.0 * d * d;
  float b = 4.0 / 29.0;

  for(long i = 0; i < npix; i++) {
    if( f[i] > d )
      I[i] = f[i]*f[i]*f[i];
    else
      I[i] = a*(f[i] - b);
    I[i] = In*I[i];
  }
}

void FitsColor::Lab_XYZ(float L, float a, float b, float *X, float *Y, float *Z)
{
  float Lx = (L + 16.0f) / 116.0f;
  float ax = Lx + a / 500.0f;
  float bx = Lx - b / 200.0f;

  *X = Xn*Lab_invfun(ax);
  *Y = Yn*Lab_invfun(Lx);
  *Z = Zn*Lab_invfun(bx);
}

float FitsColor::Lab_invfun(float f)
{
  float d = 6.0 / 29.0;
  float a = 3.0 * d * d;
  float b = 4.0 / 29.0;

  float I;
  if( f > d )
    I = f*f*f;
  else
    I = a*(f - b);

  return I;
}

void FitsColor::TuneColors(long npix, float *L, float *a, float *b)
{
  for(long i = 0; i < npix; i++) {
    float h = atan2f(b[i],a[i]) + hue;
    float r2 = a[i]*a[i] + b[i]*b[i];
    float ch = r2 > 0 ? saturation*sqrtf(r2) : 0;
    a[i] = ch*cosf(h);
    b[i] = ch*sinf(h);
  }
}

void FitsColor::NiteVision(long npix, float *L, float *X, float *Y, float *Z)
{
  if( nitevision ) {

    float xw = Xn / 100.0;
    float zw = Zn / 100.0;

    for(long i = 0; i < npix; i++) {

      // simplified transformation for photo-,meso-, and scotopic
      float r = (Y[i] - nitethresh) / nitewidth;
      float w = (1 + erf(r)) / 2;
      float s = 0.36169f*Z[i] + 1.18214f*Y[i] - 0.80498f*X[i];
      float ws = s*(1.0f - w);

      X[i] = w*X[i] + ws * xw;
      Y[i] = w*Y[i] + ws;
      Z[i] = w*Z[i] + ws * zw;

    }
  }
}
