//
// Copyright 2003 (c) Kevin Atkinson under the GNU GPL license version
// 2.0.  You should have received a copy of the GPL license along with
// this program if you did not you can find it at http://www.gnu.org/.
//

#include <vector>
#include <cassert>
#include <iostream>
#include <algorithm>

#include <yuv4mpeg.h>

using namespace std;

static inline bool is_odd(int x) {return x & 1;}
static inline bool is_even(int x) {return !is_odd(x);}

void error(const char * err, const string & p = "")
{
  cerr << "Error: " << err;
  if (!p.empty())
    cerr << p;
  cerr << "\n";
  exit(1);
}

typedef unsigned char byte;

struct IFrame
{
  int num;
  byte * even[3];
  byte * odd[3];
  float diff1;
  float diff2;
  int res1;
  int res2;
  y4m_frame_info_t inf;
  void init(size_t, size_t);
  IFrame::IFrame() {num = -1; even[0] = 0;}
  ~IFrame() {if (even[0]) free(even[0]);}
};

struct PFrame
{
  byte * d[3];
  y4m_frame_info_t inf;
  void init(size_t, size_t);
  PFrame::PFrame() {d[0] = 0;}
  ~PFrame() {if (d[0]) free(d[0]);}
};

enum DeinterlaceMethod {WEAVE, USE_EVEN, USE_ODD};

struct FrameInfo {
  // -2: use none, -1: use prev, 0: use cur, 1: use next
  int even;   // even field
  int odd;    // odd  field
  DeinterlaceMethod method;
  FrameInfo() {}
  FrameInfo(int e, int o, DeinterlaceMethod m)
    : even(e), odd(o), method(m) {}
};

//////////////////////////////////////////////////////////////////////
//
//
//

struct Analysis
{
  float thres1;
  float thres2;
  float vals[3];
  int last_set;
  int ready;

  void init();
  int adv(float val);
};

struct Compare
{
  size_t width, height;

  void subsample_line(byte * x, short * o);
  float operator() (byte * top, byte * bottom);
};

enum FieldOrder {FO_Unknown, FO_TFF, FO_BFF};

struct Match {

  y4m_stream_info_t s_stream_info;
  size_t width, c_width, height, c_height;
  size_t y_size, c_size, f_size;
  y4m_ratio_t framerate;
  FieldOrder field_order;
  bool swap_fields;

  Analysis analysis;
  Compare compare;

  IFrame * f1;
  IFrame * f2;

  int finished;

  Match() : field_order(FO_Unknown), swap_fields(false) {}

  void prep();
  void init(IFrame * f3, IFrame * f4);
  void init_iframe(IFrame *);
  IFrame * adv(IFrame *);

};


class Fields
{
public:
  Match * match;

  IFrame * f[3];

  int offset;
  int prev;      // diffrence from cur
  int from_prev; 
  int cur;

  FieldOrder field_order;

  void init(Match *, IFrame *, IFrame *, IFrame *);
  bool adv_to(int field);
  inline int field_diff(int d);
  int diff (int d1, int d2);
  bool at_end() {return f[2]->num == -1;}
};

class Chose
{
public:
  Fields fd;
  float ratio; // the decimate ratio

  float accum; // a accumulator used to insure the right number of
                // frames and to try to maintain a consistent decimate
                // pattern

  float round_div;
  
  int field;

  bool aggressive; 
  // Always match a field up with a frame and also allow a field to be
  // in a frame out of order when deinterlaceing needs to be done.
  // Will lead to better results when the material is trully 3:2
  // Pulldown (or the like).  Might cause strange results when the
  // material is trully interlaced and deinterlacing is done by
  // blending fields together.  However, if interpolation is used on
  // the moving areas of an image then this option is always safe to
  // use and should lead to better results.

  bool first;

  Chose() : accum(0), round_div(1.0), field(0), aggressive(true), first(true) {}

  void init() {}

  bool adv_field();
  FrameInfo get_frame();

  bool next(FrameInfo & f) {
    if (!first) {
      bool res = adv_field();
      if (!res) return false;
    }
    f = get_frame();
    first = false;
    return true;
  }

  bool odd(int x) {return x & 1;}
  bool even(int x) {return !odd(x);}
  bool same_group(int x, int y) {return fd.diff(x,y) <= 0;}
  bool similar_group(int x, int y) {return fd.diff(x,y) <= 1;}
  bool diff_group(int x, int y) {return fd.diff(x,y) == 2;}
};

struct Render
{
  y4m_stream_info_t d_stream_info;
  y4m_ratio_t framerate;

  size_t width, height;
  size_t c_width, c_height;

  IFrame * * f;
  PFrame out;

  bool always_weave;

  Render() : always_weave(false) {}

  void init(Match *, IFrame * *);
  void operator() (FrameInfo & finf);
  void decomb(DeinterlaceMethod, byte * even[3], byte * odd[3], byte * dest[3]);
};


//////////////////////////////////////////////////////////////////////
//
//
//

ostream & operator<< (ostream & o, const FrameInfo & finf)
{
  o << '[' << finf.even;
  if (finf.method == USE_EVEN) o << "*";
  o << ", " << finf.odd;
  if (finf.method == USE_ODD)  o << "*";
  o << "]";
  return o;
}

int main(int argc, const char *argv[])
{
   IFrame buf[5];
   Match match;
   Chose chose;
   Render render;

   for (int i = 1; i < argc; ++i)
   {
     if      (strcmp(argv[i], "tff")==0)
       match.field_order = FO_TFF;
     else if (strcmp(argv[i], "bff")==0)
       match.field_order = FO_BFF;
     else if (strcmp(argv[i], "swap")==0)
       match.swap_fields = true;
     else if (strcmp(argv[i], "normal")==0)
       chose.aggressive = false;
     else if (strcmp(argv[i], "weave")==0)
       render.always_weave = true;
     else
       error("Usage: ivtc [ttf|bff] [swap] [normal] [weave] < in.yuv > out.yuv");
   }

   match.prep();
   for (int i = 0; i != 5; ++i) match.init_iframe(&buf[i]);
   match.init(&buf[0], &buf[1]);
   chose.ratio = 60.0/24.0;
   chose.init();
   chose.fd.init(&match, &buf[2], &buf[3], &buf[4]);
   IFrame * * f = chose.fd.f;
   render.init(&match, f);

   FrameInfo finf;
   while (chose.next(finf)) 
   {
     cerr << f[1]->num << ' ' << finf << '\n';
     render(finf);
   }
}

//////////////////////////////////////////////////////////////////////
//
// Frame
//

void IFrame::init(size_t y_size, size_t c_size)
{
  even[0] = (byte *)malloc(y_size + c_size * 2);
  even[1] = even[0] + y_size/2;
  even[2] = even[1] + c_size/2;
  odd[0] = even[2] + c_size/2;
  odd[1] = odd[0] + y_size/2;
  odd[2] = odd[1] + c_size/2;
  num = -1;
}

void PFrame::init(size_t y_size, size_t c_size)
{
  d[0] = (byte *)malloc(y_size + c_size * 2);
  d[1] = d[0] + y_size;
  d[2] = d[1] + c_size;
  y4m_init_frame_info(&inf);
}

//////////////////////////////////////////////////////////////////////
//
// Compare
//

void Match::prep()
{
  y4m_init_stream_info(&s_stream_info);
  if (y4m_read_stream_header (0, &s_stream_info) != 0)
    error("Unable to read stream header from stdin");

  width = y4m_si_get_width (&s_stream_info);
  height = y4m_si_get_height (&s_stream_info);
  framerate = y4m_si_get_framerate (&s_stream_info);
  
  int interlace = y4m_si_get_interlace (&s_stream_info);

  if (field_order == FO_Unknown) {
    if (interlace == Y4M_ILACE_TOP_FIRST)
      field_order = FO_TFF;
    else if (interlace == Y4M_ILACE_BOTTOM_FIRST)
      field_order = FO_BFF;
    else
      error("Can not determine field order.");
  }

  y_size = width*height;

  c_height = height/2;
  c_width = width/2;
  c_size = c_width*c_height;

  f_size = y_size + c_size * 2;

  analysis.init();
  compare.width = width;
  compare.height = height;
}

void Match::init_iframe(IFrame * f)
{
  f->init(y_size, c_size);
}

void Match::init(IFrame * f3, IFrame * f4)
{
  f1 = f3;
  f2 = f3;
  f2->num = 0;
  if (swap_fields)
    y4m_read_fields(0, &s_stream_info, &f2->inf, f2->odd, f2->even);
  else
    y4m_read_fields(0, &s_stream_info, &f2->inf, f2->even, f2->odd);
  finished = 2;
  adv(f4);
}

IFrame * Match::adv(IFrame * f3) 
{
  if (finished == 0) {
    f3->num = -1;
    return f3;
  }

  IFrame * f0 = f1;
  f1 = f2;
  f2 = f3;
  f2->num = f1->num + 1;

  if (finished == 2) {

    int res;
    if (swap_fields)
      res = y4m_read_fields(0, &s_stream_info, &f2->inf, f2->odd, f2->even);
    else
      res = y4m_read_fields(0, &s_stream_info, &f2->inf, f2->even, f2->odd);
    
    if (res == 0) {
      f1->diff1 = compare(f1->even[0], f1->odd[0]);
      if (field_order == FO_TFF)
        f1->diff2 = compare(f2->even[0], f1->odd[0]);
      else
        f1->diff2 = compare(f1->even[0], f2->odd[0]);
    } else {
      finished--;
    }

    f0->res1 = analysis.adv(f1->diff1);
    f0->res2 = analysis.adv(f1->diff2);

  } else {

    f0->res1 = analysis.adv(0);
    f0->res2 = analysis.adv(0);

    finished--;

  }

  cerr << "Match " << f0->num << ": "
       << f0->res1 << ' ' << f0->diff1 << " | " 
       << f0->res2 << ' ' << f0->diff2 << '\n';

  return f0;
}

//////////////////////////////////////////////////////////////////////
//
// Compare
//

static const int VSCALE = 2;
static const int HSCALE = 3;

void Compare::subsample_line(byte * x, short * o)
{
  if (VSCALE == 2) {

    byte * stop = x + width;
    *o = x[0] + x[1];
    x += 2;
    ++o;
    while (x < stop) {
      *o = x[-1] + x[0] + x[1];
      x += 2;
      ++o;
    }

  } else if (VSCALE == 4) {

    byte * stop = x + width;
    *o = x[0] + x[1] + x[2] + x[3];
    x += 4;
    ++o;
    while (x < stop) {
      *o = x[-3] + x[-2] + x[-1] + x[0] + x[1] + x[2] + x[3];
      x += 4;
      ++o;
    }

  } else {

    abort();

  }
}

float Compare::operator() (byte * top, byte * bottom)
{
  int line = 0;
  int i;
  byte * s_l[2] = {top, bottom};
  int s_height = height/HSCALE;
  int s_width  = width/VSCALE;
  short ad[s_width],bd[s_width],cd[s_width],dd[s_width],ed[s_width];
  short * a = ad, * b = bd, * c = cd, * d = dd, * e = ed;
  subsample_line(top, d);
  subsample_line(bottom, e);
  float accum = 0;
  while (line < s_height)
  {
    swap(a,d);
    swap(b,e);
    i = line & 1; // 0 if even, 1 if odd
    s_l[i] += width; subsample_line(s_l[i],c);
    s_l[i] += width; subsample_line(s_l[i],e);
    i = i ^ 1; // 1 if even, 0 if odd
    s_l[i] += width; subsample_line(s_l[i],d);

    // now compare lines
    for (int x = 0; x < s_width; ++x) 
    {
      int v = a[x] + 4*c[x] + e[x] - 3*(b[x] + d[x]);
      v *= v;
      accum += v;
    }
    ++line;
  }
  static const float scale = 255* 6 * (2*VSCALE - 1);
  accum /= s_width*s_height*scale*scale;
  return accum;
}

//////////////////////////////////////////////////////////////////////
//
// Analysis
//

static const float THRES1_DEF = 0.00010;
static const float THRES2_DEF = 0.00025;

static const int MAX_LAST_SET = 20;
static const float T1_ABS = 0.0025, T1_REL = 0.65, T1_DIFF = 0.07;
static const float T1_T1 = 0.05, T1_T2 = 0.15;
static const float T2_ABS = 0.0040, T2_REL = 0.92, T2_DIFF = 0.10;
static const float T2_T2 = 0.03, T2_T1 = -0.10;

void Analysis::init()
{
  thres1 = THRES1_DEF;
  thres2 = THRES2_DEF;
  last_set = -1;
  ready = 0;
}

int Analysis::adv(float val)
{
  if (ready < 2) {

    vals[ready] = val;
    ready++;
    return -1;

  } else {

    vals[2] = val;
    if (vals[0] < vals[1] && vals[2] < vals[1]) {
      float p = vals[0] / vals[1];
      float n = vals[2] / vals[1];
      //cerr << " " << vals[0] << " " << p 
      //     << " " << abs(p - n)
      //     << " " << n << " " << vals[2] << "\n";
      if (vals[0] < T1_ABS && vals[2] < T1_ABS &&
          p < T1_REL && n < T1_REL && abs(p - n) < T1_DIFF) 
      {
        //cerr << " T1\n";
        last_set = 0;
        thres1 = (max(p,n) + T1_T1)*vals[1];
        thres2 = (max(p,n) + T1_T2)*vals[1];
      } 
      else if (vals[0] < T2_ABS && vals[2] < T2_ABS &&
               p < T2_REL && n < T2_REL && abs(p - n) < T2_DIFF) 
      {
        //cerr << " T2\n";
        last_set = 0;
        thres2 = (max(p,n) + T2_T2)*vals[1];
        if (thres1 > thres2) thres1 = (min(p,n) + T2_T1)*vals[1];
      }
    }
    if (last_set > MAX_LAST_SET) 
    {
      //cerr << " RESET\n";
      thres1 = THRES1_DEF;
      thres2 = THRES2_DEF;
      last_set = -1;
    }
    if (last_set >= 0)
      last_set ++;
    int what;
    if (vals[0] < thres1)       what = 0;
    else if (vals[0] < thres2)  what = 1;
    else                        what = 2;
    //cerr << " " << vals[0] << " = " << what << '\n';
    vals[0] = vals[1];
    vals[1] = vals[2];
    return what;
  }
}

//////////////////////////////////////////////////////////////////////
//
// Chose
//

static const float DELTA = 0.001;

bool Chose::adv_field()
{
  float ideal = field + accum + ratio;
  int     cur = (int)ideal;
  bool res = fd.adv_to(cur);
  if (!res) return false;
  int p = field - cur;  // note p does not nessary have to equal -1
  float frac  = ideal - cur;
  enum Chosen {Either, Cur, Next};
  int chosen = Cur;

  if (first || fd.at_end())                               chosen = Cur;
  else if (same_group(0, 1))                              chosen = Either;
  else if (same_group(p, 0))                              chosen = Next;
  else if (similar_group(p,0) && similar_group(-1,0)) {
    if (similar_group(0, 1) && !same_group(1, 2))         chosen = Either;
    else if (diff_group(0, 1))                            chosen = Next;
  } else if (diff_group(p, 0) && diff_group(-1, 0)
             && diff_group(1,2))                          chosen = Either;

  if (chosen == Either) {
    field = frac < round_div ? cur : cur + 1;
  } else if (chosen == Cur) {
    if (round_div < frac) round_div = frac + DELTA;
    field = cur;
  } else if (chosen == Next) {
    if (round_div > frac) round_div = frac - DELTA;
    field = cur + 1;
  }

  accum = ideal - field;
  cerr << "Chose: " << field << "  Ideal: " << ideal 
       << "  CS: " << chosen << "\n";

  return true;
}

FrameInfo Chose::get_frame()
{
  int f = field;
  bool res = fd.adv_to(f);
  assert(res);
  int matching = fd.offset == 0 ?  1 : -1;
  int other    = -matching;

  enum UseF {Matching, Other, Neither} use_f;
  bool weave = true;
  
  if      (same_group(0, matching))   use_f = Matching;
  else if (same_group(0, other))      use_f = Other;
  else if (aggressive) {
    weave = false;
    if (similar_group(0, matching))   use_f = Matching;
    else if (similar_group(0, other)) use_f = Other;
    else                              use_f = Matching;
  } else {
    weave = false;
    int other_n = -(matching + 1);
    if (!same_group(other, other_n))  use_f = Matching;
    else                              use_f = Neither;
  }

  if (fd.field_order == FO_TFF) {
    if (fd.offset == 0)
      return FrameInfo(0, use_f == Matching ? 0 : use_f == Other ? -1 : -2,
                       weave ? WEAVE : USE_EVEN);
    else
      return FrameInfo(use_f == Matching ? 0 : use_f == Other ?  1 : -2, 0,
                       weave ? WEAVE : USE_ODD);
  } else {
    if (fd.offset == 0)
      return FrameInfo(use_f == Matching ? 0 : use_f == Other ? -1 : -2, 0,
                       weave ? WEAVE : USE_ODD);
    else
      return FrameInfo(0, use_f == Matching ? 0 : use_f == Other ?  1 : -2,
                       weave ? WEAVE : USE_EVEN);
  }
}

//////////////////////////////////////////////////////////////////////
//
// Fields
//

void Fields::init(Match * m, IFrame * f0, IFrame * f1, IFrame * f2)
{
  match = m;
  f[0] = f0;
  f[1] = f1;
  f[2] = f2;

  cur = 0;

  IFrame * c = match->adv(f[0]);
  f[0] = f[1];
  f[1] = f[2];
  f[2] = c;

  field_order = m->field_order;
}

bool Fields::adv_to(int field)
{
  from_prev = 0;
  int p_fnum = cur / 2;
  int p_offset = cur % 2;
  prev = cur - field + 2;
  cur = field;
  int fnum = cur/2;
  offset = cur % 2;
  IFrame * c;
  assert(f[1]->num <= fnum);
  while ( f[1]->num != fnum && f[2]->num != -1)
  {
    if (p_fnum <= f[0]->num) {
      if (f[0]->num == p_fnum && p_offset == 1)
        from_prev = f[0]->res1;
      else
        from_prev = max(from_prev, max(f[0]->res1, f[0]->res2));}
    c = match->adv(f[0]);
    f[0] = f[1];
    f[1] = f[2];
    f[2] = c;
  }
  return f[1]->num == fnum;
}

// return the diff between the paramer frame and the next
inline int Fields::field_diff(int d)
{
  d += offset + 2;
  if (d < 0) {assert(d == prev); return from_prev;}
  IFrame * c = f[d/2];
  if (c->num == -1) return 2;
  if (is_even(d)) return c->res1;
  else            return c->res2;
}

int Fields::diff(int d1, int d2)
{
  if (d1 == d2) return 0;
  if (d1 > d2) swap(d1, d2);
  int m = 0;
  for (int i = d1; i != d2; ++i)
    m = max(m, field_diff(i));
  return m;
}

//////////////////////////////////////////////////////////////////////
//
// Render
//

void Render::init(Match * m, IFrame * * f0)
{
  f = f0 + 1;
  width = m->width; height = m->height;
  c_width = m->c_width; c_height = m->c_height;
  y4m_copy_stream_info(&d_stream_info, &m->s_stream_info);
  out.init(m->y_size, m->c_size);
  y4m_si_set_interlace(&d_stream_info, Y4M_ILACE_NONE);

  framerate = y4m_fps_NTSC_FILM;
  y4m_si_set_framerate(&d_stream_info, framerate);

  y4m_write_stream_header(1, &d_stream_info);
}

void Render::operator() (FrameInfo & finf)
{
  byte * * e = finf.even == -2 ? 0 : f[finf.even]->even;
  byte * * o = finf.odd  == -2 ? 0 : f[finf.odd ]->odd;
  decomb(finf.method, e, o, out.d);
  y4m_write_frame(1, &d_stream_info, &out.inf, out.d);
}

void Render::decomb(DeinterlaceMethod dim, 
                    byte * even[3], byte * odd[3], byte * dest[3])
{
  if (always_weave && even && odd)
    dim = WEAVE;

  if (dim == WEAVE) 
  {
    byte * dl = dest[0], * el = even[0], * ol = odd[0];

    for (size_t y = 0; y != height/2; ++y) {
      memcpy(dl, el, width); dl += width; el += width; 
      memcpy(dl, ol, width); dl += width; ol += width;}

    for (int i = 1; i != 3; ++i) {
      byte * d0 = dest[i], * d1 = dest[i] + c_width;
      el = even[i]; ol = odd[i];
      for (size_t y = 0; y != c_height/2; ++y) {
        for (size_t x = 0; x != c_width; ++x) {
          d0[x] = (el[x] + ol[x])/2;
          d1[x] = d0[x];}
        d0 += 2*c_width; d1 += 2*c_width;
        el += c_width; ol += c_width;}}
  }
  else
  {
    byte * d0 = dest[0], * d1 = dest[0] + width;

    if (dim == USE_EVEN) {

      byte * s0 = even[0]; byte * s1 = even[0] + width;
      for (size_t y = 0; y != height/2 - 2; ++y) {
        for (size_t x = 0; x != width; ++x) {
          d0[x] = s0[x];
          d1[x] = (s0[x] + s1[x])/2;}
        d0 += 2*width; d1 += 2*width;
        s0 = s1; s1 += width;}
      memcpy(d0, s0, width);
      memcpy(d1, s0, width);

    } else {
      
      byte * s0 = 0, * s1 = odd[0];
      memcpy(d0, s1, width);
      memcpy(d1, s1, width);
      for (size_t y = 1; y != height/2; ++y) {
        d0 += 2*width; d1 += 2*width;
        s0 = s1; s1 += width;
        for (size_t x = 0; x != width; ++x) {
          d0[x] = (s0[x] + s1[x])/2;
          d1[x] = s1[x];}}
    }
    
    for (size_t i = 1; i != 3; ++i) {
      byte * d = dest[i];
      byte * s = dim == USE_EVEN ? even[i] : odd[i];
      for (size_t y = 0; y != c_height/2; ++y) {
        memcpy(d, s, c_width); d += c_width;
        memcpy(d, s, c_width); d += c_width;
        s += c_width;}}
  }
}