/* MAKE-LDPC.C - Make a Low Density Parity Check code's parity check matrix. */

/* Copyright (c) 2000, 2001 by Radford M. Neal 
 *
 * Permission is granted for anyone to copy, use, or modify this program 
 * for purposes of research or education, provided this copyright notice 
 * is retained, and note is made of any changes that have been made. 
 *
 * This program is distributed without any warranty, express or implied.
 * As this program was written for research purposes only, it has not been
 * tested to the degree that would be advisable in any important application.
 * All use of this program is entirely at the user's own risk.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include "rand.h"
#include "alloc.h"
#include "intio.h"
#include "open.h"
#include "mod2sparse.h"
#include "mod2dense.h"
#include "mod2convert.h"
#include "rcode.h"


/* METHODS FOR CONSTRUCTING CODES. */

typedef enum 
{ Evencol, 	/* Uniform number of bits per column, with number specified */
  Evenboth 	/* Uniform (as possible) over both columns and rows */
} make_method; 


void make_ldpc(int,make_method,int,int);
void usage(void);


/* MAIN PROGRAM. */

int main
( int argc,
  char **argv
)
{
  make_method method;
  char *file, **meth;
  int seed, no4cycle;
  char junk;
  FILE *f;
  int cb;

  /* Look at initial arguments. */

  if (!(file = argv[1])
   || !argv[2] || sscanf(argv[2],"%d%c",&M,&junk)!=1 || M<=0
   || !argv[3] || sscanf(argv[3],"%d%c",&N,&junk)!=1 || N<=0
   || !argv[4] || sscanf(argv[4],"%d%c",&seed,&junk)!=1)
  { usage();
  }

  /* Look at the arguments specifying the method for producing the code. */

  meth = argv+5;

  if (!meth[0]) usage();

  no4cycle = 0;

  if (strcmp(meth[0],"evencol")==0 || strcmp(meth[0],"evenboth")==0)
  { method = strcmp(meth[0],"evencol")==0 ? Evencol : Evenboth;
    if (!meth[1] || sscanf(meth[1],"%d%c",&cb,&junk)!=1 || cb<=0)
    { usage();
    }
    if (meth[2])
    { if (strcmp(meth[2],"no4cycle")==0)
      { no4cycle = 1;
        if (meth[3])
        { usage();
        }
      }
      else
      { usage();
      }
    }
  }
  else
  { usage();
  }

  /* Check for some problems. */
    
  if (cb>M)
  { fprintf(stderr,
      "Number of checks per bit (%d) is greater than total checks (%d)\n",
      cb, M);
    exit(1);
  }
  
  if (cb==M && N>1 && no4cycle)
  { fprintf(stderr,
     "Can't eliminate cycles of length four with this many checks per bit\n");
    exit(1);
  }

  /* Make the parity check matrix. */

  make_ldpc(seed,method,cb,no4cycle);

  /* Write out the parity check matrix. */

  f = open_file_std(file,"wb");
  if (f==NULL) 
  { fprintf(stderr,"Can't create parity check file: %s\n",file);
    exit(1);
  }

  intio_write(f,('P'<<8)+0x80);
  
  if (ferror(f) || !mod2sparse_write(f,H) || fclose(f)!=0)
  { fprintf(stderr,"Error writing to parity check file %s\n",file);
    exit(1);
  }

  return 0;
}


/* PRINT USAGE MESSAGE AND EXIT. */

void usage(void)
{ fprintf(stderr,"Usage:  make-ldpc pchk-file n-checks n-bits seed method\n");
  fprintf(stderr,"Method: evencol  checks-per-col [ \"no4cycle\" ]\n");
  fprintf(stderr,"    or: evenboth checks-per-col [ \"no4cycle\" ]\n");
  exit(1);
}


/* CREATE A SPARSE PARITY-CHECK MATRIX.  Of size M by N, stored in H. */

void make_ldpc
( int seed,		/* Random number seed */
  make_method method,	/* How to make it */
  int cb,		/* Checks per bit */
  int no4cycle		/* Eliminate cycles of length four? */
)
{
  mod2entry *e, *f, *g, *h;
  int added, uneven, elim4;
  int i, j, k, t;
  int *u;

  rand_seed(10*seed+1);

  H = mod2sparse_allocate(M,N);

  /* Create the initial version of the parity check matrix. */

  switch (method)
  { 
    case Evencol:
    { 
      for (j = 0; j<N; j++)
      { for (k = 0; k<cb; k++)
        { do
          { i = rand_int(M);
          } while (mod2sparse_find(H,i,j));
          mod2sparse_insert(H,i,j);
        }
      }

      break;
    }

    case Evenboth:
    {
      u = chk_alloc (cb*N, sizeof *u);

      for (k = cb*N-1; k>=0; k--)
      { u[k] = k%M;
      }

      uneven = 0;
      t = 0;

      for (j = 0; j<N; j++)
      { for (k = 0; k<cb; k++)
        { 
          for (i = t; i<cb*N && mod2sparse_find(H,u[i],j); i++) ;

          if (i==cb*N)
          { uneven += 1;
            do
            { i = rand_int(M);
            } while (mod2sparse_find(H,i,j));
            mod2sparse_insert(H,i,j);
          }
          else
          { do
            { i = t + rand_int(cb*N-t);
            } while (mod2sparse_find(H,u[i],j));
            mod2sparse_insert(H,u[i],j);
            u[i] = u[t];
            t += 1;
          }
        }
      }

      if (uneven>0)
      { fprintf(stderr,"Had to place %d checks in rows unevenly\n",uneven);
      }

      break;
    }

    default: abort();
  }

  /* Add extra bits to avoid rows with less than two checks. */

  added = 0;

  for (i = 0; i<M; i++)
  { e = mod2sparse_first_in_row(H,i);
    if (mod2sparse_at_end(e))
    { j = rand_int(N);
      e = mod2sparse_insert(H,i,j);
      added += 1;
    }
    e = mod2sparse_first_in_row(H,i);
    if (mod2sparse_at_end(mod2sparse_next_in_row(e)) && N>1)
    { do 
      { j = rand_int(N); 
      } while (j==mod2sparse_col(e));
      mod2sparse_insert(H,i,j);
      added += 1;
    }
  }

  if (added>0)
  { fprintf(stderr,
           "Added %d extra bit-checks to make row counts at least two\n",
           added);
  }

  /* Add extra bits to try to avoid problems with even column counts. */

  if (cb%2==0 && cb<M && N>1 && added<2)
  { int a;
    for (a = 0; added+a<2; a++)
    { do
      { i = rand_int(M);
        j = rand_int(N);
      } while (mod2sparse_find(H,i,j));
      mod2sparse_insert(H,i,j);
    }
    fprintf(stderr,
 "Added %d extra bit-checks to try to avoid problems from even column counts\n",
      a);
  }

  /* Eliminate cycles of length four, if asked, and if possible. */

  if (no4cycle)
  { 
    elim4 = 0;

    for (t = 0; t<10; t++) 
    { k = 0;
      for (j = 0; j<N; j++)
      { for (e = mod2sparse_first_in_col(H,j);
             !mod2sparse_at_end(e);
             e = mod2sparse_next_in_col(e))
        { for (f = mod2sparse_first_in_row(H,mod2sparse_row(e));
               !mod2sparse_at_end(f);
               f = mod2sparse_next_in_row(f))
          { if (f==e) continue;
            for (g = mod2sparse_first_in_col(H,mod2sparse_col(f));
                 !mod2sparse_at_end(g);
                 g = mod2sparse_next_in_col(g))
            { if (g==f) continue;
              for (h = mod2sparse_first_in_row(H,mod2sparse_row(g));
                   !mod2sparse_at_end(h);
                   h = mod2sparse_next_in_row(h))
              { if (mod2sparse_col(h)==j)
                { do
                  { i = rand_int(M);
                  } while (mod2sparse_find(H,i,j));
                  mod2sparse_delete(H,e);
                  mod2sparse_insert(H,i,j);
                  elim4 += 1;
                  k += 1;
                  goto nextj;
                }
              }
            }
          }
        }
      nextj: ;
      }
      if (k==0) break;
    }

    if (elim4>0)
    { fprintf(stderr,
        "Eliminated %d cycles of length four by moving checks within column\n",
         elim4);
    }

    if (t==10) 
    { fprintf(stderr,
        "Couldn't eliminate all cycles of length four in 10 passes\n");
    }
  }
}