/*

  xmunipack - color space

  Copyright © 2009-2014 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>



// ---- 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),nvision(false),saturation(1.0),hue(0.0),nthresh(0.0),
  nwidth(30.0)
{
  UnRef();
  SetRefData(new FitsColorData);

  InitUV();
}

FitsColor::FitsColor(const wxString& cdatafile, const FitsArray& array):
  ispace(COLORSPACE_XYZ),nvision(false),saturation(1.0),hue(0.0),nthresh(0.0),
  nwidth(30.0)
{
  UnRef();
  SetRefData(new FitsColorData);

  InitUV();
  Init(cdatafile,array);
}

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

void FitsColor::InitUV()
{
  uwhite = 0.211;
  vwhite = 0.475;

  // white point for D65 and 2 deg. observer
  uwhite = 0.2009;
  vwhite = 0.4610;

  // empiricaly determined by minimizing of saturation for a gray image
  uwhite = 0.19784;
  vwhite = 0.46831;
} 

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::SetWhitePoint(float u, float v)
{
  if( u >= 0.0 )
    uwhite = u;
  if( v >= 0.0 )
    vwhite = v;
}

void FitsColor::SetNightThresh(float x)
{
  nthresh = x;
}

void FitsColor::SetNightWidth(float x)
{
  nwidth = x;
}

void FitsColor::SetNightVision(bool t)
{
  nvision = t;
}

float FitsColor::NightProfile(float x) const
{
  return 1.0/(1.0 + expf(-2.5*(x - nthresh)/nwidth));
}

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::ConvertRGB(float X, float Y, float Z, 
			   unsigned char &R, unsigned char &G, unsigned char &B)
{
  float L,u,v;
  XYZ_Luv(X,Y,Z,Yn,L,u,v);

  // color adjustment
  float h = (v != 0.0 && u != 0.0 ? atan2f(v,u) : 0.0) + hue;
  float ch = satur*sqrtf(v*v + u*u);
  u = ch*cosf(h);
  v = ch*sinf(h);
  
  if( nvision ) {
    float w = NightProfile(L);
    float s = Scotopic(X,Y,Z);
    u = w*u;
    v = w*v;
    L = w*L + (1.0 - w)*s;
  }

  Luv_XYZ(L,u,v,1.0,X,Y,Z);

  OutputRGB(X,Y,Z,R,G,B);
}
*/

 /*
void FitsColor::Transform(size_t n, float *t, float& Z, float& Y, float& X)
{
  float *c[] = { &Z, &Y, &X };
  int ncolors = GetColors();
  int nbands = GetBands();
  for(int i = 0; i < ncolors; i++) {
    float s = 0.0;
    for(int j = 0; j < nbands; j++)
      s = s + GetTrans(j,i)*t[j];
    **(c+i) = s;
  }
  return;


  float a[3][3] = {{4.961444,   0.096640,  -0.021961},{0.293121,   1.479324,   0.348571},{0.902866,   0.513487,   0.837347}};
  float b[3];

  for(size_t i = 0; i < 3; i++) {
    float s = 0.0;
    for(size_t j = 0; j < n; j++)
      s = s + a[i][j]*t[j];
    b[i] = s;
  }

  Z = b[0];
  Y = b[1];
  X = b[2];
}
 */

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

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

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

    int 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(int 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];
    }
  }

}



// http://en.wikipedia.org/wiki/CIELUV_color_space
void FitsColor::XYZ_Luv(int npix, float *X, float *Y, float *Z, float Yn,
			float *L, float *u, float *v)
{
  wxASSERT( Yn > 0.0);

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

    float s = X[i] + 15.0f*Y[i] + 3.0f*Z[i];
    float u1 = 0.0;
    float v1 = 0.0;
    if( s > 0.0f ) {
      u1 = 4.0f*X[i]/s;
      v1 = 9.0f*Y[i]/s;
    }

    float yy = Y[i]/Yn;
    if( yy > 0.0088565f /* = powf(6.0/29.0,3)*/ )
      //      L[i] = 116.0f*powf(yy,0.3333333f) - 16.0f;
      L[i] = 116.0f*cbrtf(yy) - 16.0f;
    else
      L[i] = 903.30f*yy;   /* = pow(26.0/3.0,3) */

    float t = 13.0f*L[i];
    u[i] = t*(u1 - uwhite);
    v[i] = t*(v1 - vwhite);
  }
}

void FitsColor::Luv_XYZ(int npix, float *L, float *u, float *v, float Yn,
			float *X, float *Y, float *Z)
{
  float y0 = Yn*0.0011071f;

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

    float s = 13.0f*L[i];
    float u1 = uwhite;
    float v1 = vwhite;
    if( s > 0.0 ) {
      u1 = u[i]/s + uwhite;
      v1 = v[i]/s + vwhite;
    }

    if( L[i] <= 8.0f ) 
      Y[i] = L[i]*y0; /* = powf(3.0/29.0,3) */
    else {

      float t = (L[i] + 16.0f)/116.0f;
      Y[i] = t*t*t*Yn;
    }

    float t = Y[i]/(4.0f*v1);
    X[i] = t*(9.0f*u1);
    Z[i] = t*(12.0f - 3.0f*u1 - 20.0f*v1);

    // ?? - whites saturated stars
    /*
    if( Y[i] > 1.0 ) Y[i] = 1.0;
    if( X[i] > 1.0 ) X[i] = 1.0;
    if( Z[i] > 1.0 ) Z[i] = 1.0;
    */
  }
}

void FitsColor::TuneColors(int npix, float *L, float *u, float *v)
{
  for(int i = 0; i < npix; i++) {

    float h = (v[i] != 0.0 && u[i] != 0.0 ? atan2f(v[i],u[i]) : 0.0) + hue;
    float ch = saturation*sqrtf(v[i]*v[i] + u[i]*u[i]);
    u[i] = ch*cosf(h);
    v[i] = ch*sinf(h);
  }
}

void FitsColor::NightVision(int npix, float *L, float *X, float *Y, float *Z)
{
  if( nvision ) {

    float t = 2.5f/nwidth;

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

      // simplified transformation for gray
      float w = 1.0f/(1.0f + expf(-t*(L[i] - nthresh)));
      float s = 0.36169f*Z[i] + 1.18214f*Y[i] - 0.80498f*X[i];
      float w1 = 1.0f - w;
      float ws = w1*s;
      X[i] = w*X[i] + ws;
      Y[i] = w*Y[i] + ws;
      Z[i] = w*Z[i] + ws;
    }
  }
}

