/* ======================================================================== *\
!
! *
! * This file is part of CheObs, the Modular Analysis and Reconstruction
! * Software. It is distributed to you in the hope that it can be a useful
! * and timesaving tool in analysing Data of imaging Cerenkov telescopes.
! * It is distributed WITHOUT ANY WARRANTY.
! *
! * Permission to use, copy, modify and distribute this software and its
! * documentation for any purpose is hereby granted without fee,
! * provided that the above copyright notice appears in all copies and
! * that both that copyright notice and this permission notice appear
! * in supporting documentation. It is provided "as is" without express
! * or implied warranty.
! *
!
!
!   Author(s): Thomas Bretz, 6/2019 <mailto:tbretz@physik.rwth-aachen.de>
!
!   Copyright: CheObs Software Development, 2000-2019
!
!
\* ======================================================================== */

//////////////////////////////////////////////////////////////////////////////
//
//  MOptics
//
//////////////////////////////////////////////////////////////////////////////
#include "MOptics.h"

#include <TRandom.h>

#include "MQuaternion.h"
#include "MReflection.h"

ClassImp(MOptics);

using namespace std;

// --------------------------------------------------------------------------
//
// Default constructor
//
MOptics::MOptics(const char *name, const char *title)
{
    fName  = name  ? name  : "MOptics";
    fTitle = title ? title : "Optics base class";
}

// --------------------------------------------------------------------------
//
// Critical angle for total reflection asin(n2/n1), or 90deg of n1<n2
// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
//
double MOptics::CriticalAngle(double n1, double n2)
{
    return n1>=n2 ? asin(n2/n1) : TMath::Pi()/2;
}

// --------------------------------------------------------------------------
//
// This returns an approximation for the reflectivity for a ray
// passing from a medium with refractive index n1 to a medium with
// refractive index n2. This approximation is only valid for similar
// values of n1 and n2 (e.g. 1 and 1.5 but not 1 and 5).   /
//
// For n1<n2 the function has to be called with theta being the
// angle between the surface-normal and the incoming ray, for
// n1>n2, alpha is the angle between the outgoing ray and the
// surface-normal.
//
// It is not valid to call the function for n1>n2 when alpha exceeds
// the critical angle. Alpha is defined in the range [0deg, 90deg]
// For Alpha [90;180], The angle is assumed  the absolute value of the cosine is used.
//
// alpha must be given in radians and defined between [0deg and 90deg]
//
// Taken from:
// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
//
double MOptics::SchlickReflectivity(double alpha, double n1, double n2)
{
    const double r0 = (n1-n2)/(n1+n2);
    const double r2 = r0 * r0;

    const double p  = 1-cos(alpha);

    return r2 + (1-r2) * p*p*p*p*p;
}

// --------------------------------------------------------------------------
//
// Same as SchlickReflectivity(double,double,double) but takes the direction
// vector of the incoming or outgoing ray as u and the normal vector of the
// surface. The normal vector points from n2 to n1.
//
double MOptics::SchlickReflectivity(const TVector3 &u, const TVector3 &n, double n1, double n2)
{
    const double r0 = (n1-n2)/(n1+n2);
    const double r2 = r0 * r0;

    const double c = -n*u;

    const double p  = 1-c;

    return r2 + (1-r2) * p*p*p*p*p;
}

// --------------------------------------------------------------------------
//
// Applies refraction to the direction vector u on a surface with the normal
// vector n for a ray coming from a medium with refractive index n1
// passing into a medium with refractive index n2. Note that the normal
// vector is defined pointing from medium n2 to medium n1 (so opposite
// of u).
//
// Note that u and n must be normalized before calling the function!
//
// Solution accoridng to (Section 6)
// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
//
// u_out = n1/n2 * u_in + (n1/n2 * cos(theta_in) - sqrt(1-sin^2(theta_out)) * n
// sin^2(theta_out) = (n1/n2)^2 * sin^2(theta_in) = (n1/n2)^2 * (1-cos^2(theta_in))
// cos(theta_in) = +- u_in * n
//
// In case of success, the vector u is altered and true is returned. In case
// of failure (incident angle above critical angle for total internal reflection)
// vector u stays untouched and false is returned.
//
bool MOptics::ApplyRefraction(TVector3 &u, const TVector3 &n, double n1, double n2)
{
    // The vector should be normalized already
    // u.NormalizeVector();
    // n.NormalizeVector();

    // Usually: Theta [0;90]
    // Here:    Theta [90;180] => c=|n*u| < 0

    const double r = n1/n2;
    const double c = -n*u;

    const double rc = r*c;

    const double s2 = r*r - rc*rc; // sin(theta_out)^2 = r^2*(1-cos(theta_in)^2)

    if (s2>1)
        return false;

    const double v = rc - sqrt(1-s2);

    u = r*u + v*n;

    return true;
}

// --------------------------------------------------------------------------
//
// In addition to calculating ApplyRefraction(u.fVectorPart, n, n1, n2),
// the fourth component (inverse of the speed) is multiplied with n1/n2
// of tramission was successfull (ApplyRefraction retruned true)
//
// Note that .fVectorPart and n must be normalized before calling the function!
//
bool MOptics::ApplyRefraction(MQuaternion &u, const TVector3 &n, double n1, double n2)
{
    if (!ApplyRefraction(u.fVectorPart, n, n1, n2))
        return false;

    u.fRealPart *= n2/n1;
    return true;
}

// --------------------------------------------------------------------------
//
// Applies a transition from one medium (n1) to another medium (n2)
// n1 is always the medium where the photon comes from.
//
// Uses Schlick's approximation for reflection
//
// The direction of the normal vector does not matter, it is automatically
// aligned (opposite) of the direction of the incoming photon.
//
// Based on
// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
//
// Returns:
//
//  0 error (total internal reflection from optical thin to optical thick medium)
//
//  1 total internal reflection applied
//  2 Monte Carlo    reflection applied
//
//  3 nothing done (n1==n2, transmission)
//  4 refraction applied from optical thick to optically thinner medium
//  5 refraction applied from optical thin  to optically thicker medium
//
int MOptics::ApplyTransitionFast(TVector3 &dir, TVector3 n, double n1, double n2)
{
    if (n1==n2)
        return 3;

    // The normal vector must point in the same direction as
    // the photon is moving and thus cos(theta)=[90;180]
    if (dir*n>0)
        n *= -1;

    TVector3 u(dir);

    // From optical thick to optical thin medium
    if (n1>n2)
    {
        // Check for refraction...
        // ... if possible:     use exit angle for calculating reflectivity
        // ... if not possible: reflect ray
        if (!ApplyRefraction(u, n, n1, n2))
        {
            // Total Internal Reflection
            dir *= MReflection(n);
            return 1;
        }
    }

    const double reflectivity = SchlickReflectivity(u, n, n1, n2);

    // ----- Case of reflection ----

    if (gRandom->Uniform()<reflectivity)
    {
        dir *= MReflection(n);
        return 2;
    }

    // ----- Case of refraction ----

    // From optical thick to optical thin medium
    if (n1>n2)
    {
        // Now we know that refraction was correctly applied
        dir = u;
        return 4;
    }

    // From optical thin to optical thick medium
    // Still need to apply refraction

    if (ApplyRefraction(dir, n, n1, n2))
        return 5;

    // ERROR => Total Internal Reflection not possible
    // This should never happen
    return 0;
}

// --------------------------------------------------------------------------
//
// Applies a transition from one medium (n1) to another medium (n2)
// n1 is always the medium where the photon comes from.
//
// Uses Fresnel's equation for calculating reflection. Total internal
// reflection above the critical angle will always take place. Fresnel
// reflection will only be calculated if 'fresnel' is set to true (default).
// For Fresnel reflection, a random number is produced according to the
// calculated reflectivity to decide whether the ray is reflected or
// refracted.
//
//
// The direction of the normal vector does not matter, it is automatically
// aligned (opposite) of the direction of the incoming photon.
//
// Based on
// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
//
// Returns:
//
//  0 error (total internal reflection from optical thin to optical thick medium)
//
//  1 total internal reflection applied
//  2 Monte Carlo    reflection applied
//
//  3 nothing done (n1==n2, transmission)
//  4 refraction applied
//
int MOptics::ApplyTransition(TVector3 &u, TVector3 n, double n1, double n2, bool fresnel)
{
    if (n1==n2)
        return 3;

    // The normal vector must point in the same direction as
    // the photon is moving and thus cos(theta)=[90;180]
    if (u*n>0)
        n *= -1;

    // Calculate refraction outgoing direction (Snell's law)
    const double r   = n1/n2;
    const double ci  = -n*u;           // cos(theta_in)
    const double rci =  r*ci;          // n1/n2 * cos(theta_in)
    const double st2 = r*r - rci*rci;  // sin(theta_out)^2 = r^2*(1-cos(theta_in)^2)

    if (st2>1)
    {
        // Total Internal Reflection
        u *= MReflection(n);
        return 1;

    }

    const double ct  = sqrt(1-st2);    // cos(theta_out) = sqrt(1 - sin(theta_out)^2)
    const double rct = r*ct;           // n1/n2 * cos(theta_out)

    // Calculate reflectivity for none polarized rays (Fresnel's equation)
    const double Rt = (rci - ct)/(rci + ct);
    const double Rp = (rct - ci)/(rct + ci);

    const double reflectivity = (Rt*Rt + Rp*Rp)/2;

    if (fresnel && gRandom->Uniform()<reflectivity)
    {
        // ----- Case of reflection ----
        u *= MReflection(n);
        return 2;
    }
    else
    {
        // ----- Case of refraction ----
        u = r*u + (rci-ct)*n;
        return 4;
    }
}

int MOptics::ApplyTransitionFast(MQuaternion &u, const TVector3 &n, double n1, double n2)
{
    const int rc = ApplyTransitionFast(u.fVectorPart, n, n1, n2);
    if (rc>=3)
        u.fRealPart *= n2/n1;
    return rc;
}

int MOptics::ApplyTransition(MQuaternion &u, const TVector3 &n, double n1, double n2, bool fresnel)
{
    const int rc = ApplyTransition(u.fVectorPart, n, n1, n2, fresnel);
    if (rc>=3)
        u.fRealPart *= n2/n1;
    return rc;
}