#
#   hw.R - A set of R functions for calculating the Hardy Weinberg
#             equilibrium for the ABO Blood type system
#
#Copyright (C) 2004 Fernando Henrique F. P. Rosa <feferraz at ime.usp.br>
#   
# 
#This program is free software; you can redistribute it and/or
#modify it under the terms of version 2 of the GNU General Public License
#as published by the Free Software Foundation.  A copy of this license should
#be included in the file COPYING.
#
#This program 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 this program; if not, write to the Free Software
#Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
#


##---->  hw.R
##	Implementação de uma classe hardy.weinberg cálculo do
##	equilíbrio de H-W para o sistema ABO
##<----


##
## O código fonte está dividido em 3 seções, de acordo com a
## finalidade de suas funções.
##
## A primeira seção são as funções internas, utilizadas 
## para cálculo dos chutes iniciais e dos processos de otimização.
## essas funções não são usadas diretamente pelo usuário.
##
## A segunda seção é a função principal, invocada pelo usuário
## para criar objetos de análise da classe hardy.weinberg.
## Ela inclui a função hardy.weinberg(), que é utilizada pelo usuário
## para rodar a análise.
##
## A terceira seção consiste das funções que implementam os métodos
## (plot, latex, print) para objetos da classe hardy.weinberg.
## Elas são usadas direta ou indiretamente pelo usuário para 
## estudar o resultado da análise.
##


######################################################################
####### Funções Internas para Cálculo das Rotinas ####################
#######               de Otimização               ####################
######################################################################

##
## a função chutes.inicias() faz os chutes inicias para beta, usando
## o estimador dos momentos.
##

chutes.iniciais <- function(dados) {
  prop <- dados/sum(dados)
  a <- 1 - sqrt(prop[1]+prop[3])
  b <- 1 - sqrt(prop[1]+prop[2])
  o <- 1 - a - b
  res <- c(a,b,o)
  names(res) <- c("A","B","O")
  res
}

##
## a função lagrange(), recebe os parametros eps.min e n.max, que são
## passados via interface com o usuário pela função hardy.weinberg(), 
## recebe os chutes.iniciais() também da hardy.weinberg() e realiza o
## o processo iterativo por Multiplicadores de Lagrange para encontrar
## a estimativa de beta.
## Ela retorna o vetor estimado beta, o numero de iteracoes efetuadas
## e o nome do método utilizado.
##

lagrange <- function(dados,chute,eps.min=0.0001,n.max=30) {
  prop <- dados/sum(dados)
  a <- prop[2]*(chute["A"] + chute["O"])/(chute["A"]+2*chute["O"]) + prop[4]/2
  b <- prop[3]*(chute["B"] + chute["O"])/(chute["B"]+2*chute["O"]) + prop[4]/2
  o <- prop[2]*chute["O"]/(chute["A"]+2*chute["O"])+1*prop[1]+ prop[3]*chute["O"]/(chute["B"]+2*chute["O"])

  n <- 1
  eps <- sqrt(sum(c(a-chute["A"],b-chute["B"],o-chute["O"])^2))

  while (eps > eps.min && n <= n.max) {
      a.novo <- prop[2]*(a + o)/(a+2*o) + prop[4]/2
      b.novo <- prop[3]*(b + o)/(b+2*o) + prop[4]/2
      o.novo <- prop[2]*o/(a+2*o)+1*prop[1]+ prop[3]*o/(b+2*o)
      eps <- sqrt(sum(c(a-a.novo,b-b.novo,o-o.novo)^2))
      n <- n+1
      a <- a.novo
      b <- b.novo
      o <- o.novo
  }
  if (n > n.max) {
     warning("Max number of iterations reached")
  }
  res <- c(a,b,o)
  names(res) <- c("A","B","O")
  list(beta=res,n=n,method="Multiplicadores de Lagrange")
}

##
## A função U(), retorna a partir de beta e n, o valor da função
## score(). Ela foi definida aqui para poder ser chamada pelas
## funções newton.raphson() e scoring.fisher(), pois a função score
## é utilizada nesses dois métodos iterativos nas suas relações
## de recorrência.
##

U <- function(beta,n) {
      bo <-  1 - sum(beta)
      ba <- beta[1]
      bb <- beta[2]

      u1 <- 2*bo*n[2]/(ba^2 + 2*ba*bo) - 2*n[3]/(bb+2*bo) + n[4]/ba - 2*n[1]/bo
      u2 <- 2*bo*n[3]/(bb^2 + 2*bb*bo) - 2*n[2]/(ba+2*bo) + n[4]/bb - 2*n[1]/bo
     matrix(c(u1,u2),ncol=1)
}

##
## A função newton.raphson() recebe os mesmos parâmetros que a 
## lagrange(), através da hardy.weinberg(), e faz o procedimento
## iterativo via Newton-Raphson. Dentro dela, define-se a função
## hessiana(), que retorna a estimativa da matriz hessiana para um
## dado vetor de betas e ns. 
## No procedimento iterativo, "%*%" indica multiplicação matricial, 
## e solve(A) calcula a inversa da matriz A.
## O processo iterativo é rodado até que se chegue a alguma das
## condições pré-estabelecidas.
## Ela retorna o vetor estimado beta, o numero de iteracoes efetuadas
## e o nome do método utilizado.
##

newton.raphson <- function(dados,chute,eps.min=0.0001,n.max=30) {
  hessiana <- function(beta,n) {
      bo <-  1 - sum(beta)
      ba <- beta[1]
      bb <- beta[2]
      h11 <- -2*(ba^2+2*ba*bo+2*bo^2)*n[2]/ (ba^2 + 2*ba*bo)^2 - 4*n[3]/(bb+2*bo)^2 - n[4]/ba^2 - 2*n[1]/bo^2
      h12 <- -2*n[2]/(ba+2*bo)^2 - 2*n[3]/(bb+2*bo)^2 - 2*n[1]/bo^2
      h22 <- -2*(bb^2+2*bb*bo+2*bo^2)*n[3]/ (bb^2 + 2*bb*bo)^2 - 4*n[2]/(ba+2*bo)^2 - n[4]/bb^2 -
2*n[1]/bo^2
      matrix(c(h11,h12,h12,h22),ncol=2)
  }

  Chute <- matrix(c(chute["A"],chute["B"]),ncol=1)
  B <- Chute - solve(hessiana(Chute,dados)) %*% U(Chute,dados)
 
  n <- 1
  eps <- sqrt(sum((B-Chute)^2))

  while (eps > eps.min && n <= n.max) {
      B.novo <- B - solve(hessiana(B,dados)) %*% U(B,dados)
      eps <- sqrt(sum((B-B.novo)^2))
      n <- n+1
      B <- B.novo
  }
  if (n > n.max) {
     warning("Max number of iterations reached")
  }
  a <- B[1]
  b <- B[2]
  o <- 1 - (a+b)
  res <- c(a,b,o)
  names(res) <- c("A","B","O")
  list(beta=res,n=n,method="Newton-Raphson")

}

##
## A função scoring.fisher() recebe os mesmos parâmetros que a 
## lagrange(), e newton.raphson()  através da hardy.weinberg(),
## e faz o procedimento iterativo via scoring de Fisher.
## Dentro dela, define-se a função inf.fisher(), que calcula a 
## estimativa da matriz de Fisher para um dado vetor de betas e ns. 
## Ela é definida dentro da scoring.fisher() porque a cada passo
## do processo iterativo se calcula a matriz de informação de Fisher
## estimada.
## O processo iterativo é rodado até que se chegue a alguma das
## condições pré-estabelecidas.
## Ela retorna o vetor estimado beta, o numero de iteracoes efetuadas
## e o nome do método utilizado.
##

scoring.fisher <-  function(dados,chute,eps.min=0.0001,n.max=30) {
   inf.fisher <- function(beta,dados) {
     N <- sum(dados)
     bo <- 1 - sum(beta)
     ba <- beta[1]
     bb <- beta[2] 
     i12 <- 2*N*((ba^2+2*ba*bo)/(ba+2*bo)^2 + (bb^2+2*bb*bo)/(bb+2*bo)^2 + 1)
     i11 <- 2*N*((ba^2 + 2*ba*bo)*(ba^2+2*ba*bo+2*bo^2)/(ba^2+2*ba*bo)^2 + 2*(bb^2 + 2*bb*bo)/(bb + 2*bo)^2 + bb/ba + 1)
     i22 <- 2*N*((bb^2 + 2*bb*bo)*(bb^2+2*bb*bo+2*bo^2)/(bb^2+2*bb*bo)^2 + 2*(ba^2 + 2*ba*bo)/(ba + 2*bo)^2 + ba/bb + 1)
     matrix(c(i11,i12,i12,i22),ncol=2)
   }

  Chute <- matrix(c(chute["A"],chute["B"]),ncol=1)
  B <- Chute + solve(inf.fisher(Chute,dados)) %*% U(Chute,dados)
  n <- 1
  eps <- sqrt(sum((B-Chute)^2))

  while (eps > eps.min && n <= n.max) {
      B.novo <- B + solve(inf.fisher(B,dados)) %*% U(B,dados)
      eps <- sqrt(sum((B-B.novo)^2))
      n <- n+1
      B <- B.novo
  }
  if (n > n.max) {
     warning("Max number of iterations reached")
  }
  B

  a <- B[1]
  b <- B[2]
  o <- 1 - (a+b)
  res <- c(a,b,o)
  names(res) <- c("A","B","O")
  list(beta=res,n=n,method="Scoring de Fisher")

}

######################################################################
####### Função Principal, usada para Interação com ###################
#######  o usuário e criação de objetos da classe  ###################
#######               hardy.weinberg               ###################
######################################################################

##
## Essa é a função principal do programa, com a qual o usuário interage
## para fazer suas análises. 
## Ela recebe como parâmetros os dados, o eps.min, n.max e o método 
## desejado. Ela calcula então os chutes iniciais, determina o método
## escolhido e chama a função interna apropriada, de acordo com o método
## escolhido. 
## A partir dai, a partir do modelo estrutural, ela calcula as estimativas
## do vetor de parâmetros teta, e o vetor de valores esperados mi.
## Calcula ainda os valores das estatísticas Qv, Qp e Qn, seus respectivos
## P-valores, e armazena isso num data.frame.
## Por fim, a função junta todas as informações passadas pelo usuário e o 
## resultado da análise, num objeto da classe hardy.weinberg, e retorna
## esse objeto para o usuário.
##

hardy.weinberg <- function(dados,eps.min=0.0001,n.max=30,method=c("lagrange","newton","fisher")) {
	
  chutes <- chutes.iniciais(dados)
  method <- match.arg(method)
  res <- switch(method,lagrange=lagrange(dados,chutes,eps.min,n.max),newton=newton.raphson(dados,chutes,eps.min,n.max),fisher=scoring.fisher(dados,chutes,eps.min,n.max))

  teta <- numeric(4)
  names(teta) <- c("O","A","B","AB")
  teta["O"] <- res$beta["O"]^2
  teta["A"] <- res$beta["A"]^2 + 2*res$beta["O"]*res$beta["A"]
  teta["B"] <- res$beta["B"]^2 + 2*res$beta["O"]*res$beta["B"] 
  teta["AB"] <- 2*res$beta["A"]*res$beta["B"]
  
  esperados <- sum(dados)*teta
  Qv <- -2*dados%*%(log(esperados) - log(dados))
  Qp <- sum((dados-esperados)^2/esperados)
  Qn <- sum((dados-esperados)^2/dados)
  estatisticas <- data.frame(obs=c(Qv,Qp,Qn),gl=c(1,1,1),"P-valor"=c(0,0,0),row.names=c("Qv","Qp","Qn"),check.names=FALSE)
  estatisticas[1,3] <- 1-pchisq(Qv,1)
  estatisticas[2,3] <- 1-pchisq(Qp,1)
  estatisticas[3,3] <- 1-pchisq(Qn,1)

  dados[5] <- sum(dados)
  names(dados) <- c("O","A","B","AB","N")
  
  retorno <- list(dados=dados,beta=res$beta,teta=teta,esperados=esperados,estatisticas=estatisticas,method=res$method,n=res$n,eps=eps.min,n.max=n.max,chutes=chutes)
  class(retorno) <- "hardy.weinberg"
  retorno
}

######################################################################
####### Funções que registram os métodos da classe ###################
#######               hardy.weinberg               ###################
######################################################################

##
## A função principal por si só só retornaria para o  usuário
## uma lista enorme de valores, dificeis de serem visualizados. 
## É através dos métodos da classe hardy.weinberg(), que o usuário
## de fato intarege com a análise e interpreta os resultados.
## O método print é aquele chamado quando o usuário digita o nome do 
## objeto no prompt interativo do R, ou quando usa o comando print(obj).
## Ele organiza as informações contidas no objeto da classe hardy.weinberg
## e imprime de maneira parcimoniosa para o usuário.
##

print.hardy.weinberg <- function(obj) {
 cat("\n    Equilíbrio de Hardy-Weinberg para o sistema ABO \n\n")
 cat("  Dados Utilizados\n")
 print(obj$dados)
 cat("\n  Método de Estimação Utilizado\n")
 cat("   ",obj$method,"\n    eps.min =",obj$eps,"    n.max =",obj$n.max,"\n")
 cat("    n efetivo =",obj$n,"\n\n")
 cat("  Chutes iniciais para beta\n")
 print(obj$chutes)
 cat("\n  Estimativas finais de beta\n")
 print(obj$beta)
 cat("\n  Estimativas finais de teta\n")
 print(obj$teta)
 cat("\n  Valores Esperados\n")
 print(round(obj$esperados))
 cat("\n  Estatisticas de Teste\n")
 print(obj$estatisticas)
 cat("\n")
 return(invisible(obj))
}

##
## O método plot se encarrega de gerar os gráficos
## quando o usuário digita plot(obj)
## Basicamente, dividimos a tela em 6 pedaços e fazemos então
## cada gráfico em seqüência, um por vez.
## A função plota.pvalor(), definida dentro do método plot, 
## se encarrega de fazer o gráfico da distribuição Qui-Quadrado
## com os graus de liberdade adequados, e de preencher a área
## do gráfico equivalente ao P-Valor.
## Isso é feito através de um polígono com muitos lados, de acordo
## com o método descrito em:
## http://www.ime.usp.br/~feferraz/br/shaded.html
##

plot.hardy.weinberg <- function(obj) {
 erase.screen()
 split.screen(c(1,2))
 split.screen(c(2,1),1)
 split.screen(c(3,1),2)
 bp.lim <- range(obj$dados[-5],obj$esperados)
 screen(3)
 barplot(obj$dados[-5],ylim=bp.lim,col=topo.colors(4),main="Valores Observados")
 screen(4)
 barplot(obj$esperados,ylim=bp.lim,col=topo.colors(4),main="Valores Esperados\nsob Hardy-Weinberg")
 plota.pvalor <- function(estatistica) { 
    par(mar=c(2,2,2,2))
    Qobs <- estatistica[[1]]
    gl <- estatistica[[2]]
    cord.x <- c(Qobs,seq(Qobs,Qobs+2,0.01),Qobs+2)
    cord.y <- c(0,dchisq(seq(Qobs,Qobs+2,0.01),gl),0)
    curve(dchisq(x,gl),xlim=c(0,Qobs+2),main=paste(row.names(estatistica),": P-valor= ",format(estatistica[[3]],digits=4),sep=''),ylab="",yaxt="n")
    polygon(cord.x,cord.y,col="skyblue",lty=0)
 }

 screen(5) 
 plota.pvalor(obj$estatistica[1,])

 screen(6)
 plota.pvalor(obj$estatistica[2,])

 screen(7)
 plota.pvalor(obj$estatistica[3,])

}

## 
## Por fim o método latex, gera o relatório
## em LaTeX, a partir de um objeto da classe
## hardy.weinberg.
## O código parece um pouco bagunçado mas é por causa
## do uso de \\ ao invés de \ dentro das strings que definem 
## a saída em LaTeX, mas isso é feito porque a função 
## cat() considera uma barra só como caracter especial, e 
## então é preciso duas barras para obter a saída de uma.
## Basicamente o que é feito é gerar a saída em LaTeX dentro de um
## bloco minipage, com largura de 90% do texto, centrado verticalmente.
## 

latex.hardy.weinberg <- function(obj,file="",append=TRUE) {
  options(digits=5)
  saida <- paste("\\begin{minipage}[c]{0.9\\textwidth}\n",
		"\\begin{center}\\large{ Equilíbrio de Hardy-Weinberg para o sistema ABO }\\end{center}\n\n",
		"Dados considerados:\n",
		"\\begin{displaymath}\n",
		" \\stackunder{n'}{\\sim} = (",format(obj$dados[1]),",",
		format(obj$dados[2]),",",format(obj$dados[3]),",",
		format(obj$dados[4]),")",
		"\n\\end{displaymath}\n",
		"Método de estimação utilizado: ",obj$method,
		". Critérios de convergência: $\\epsilon_{\\min} = ",
		format(obj$eps),"$, $n_{\\max} = ",format(obj$n.max),
		"$. Número de iterações realizadas: $n = ",format(obj$n),"$.\n",
		"\n\\begin{displaymath}\\stackunder{\\beta^{(0)}}{\\sim} =",
		"\\left(\\begin{array}{l}\n",format(obj$chutes[1]),"\\\\\n",
		format(obj$chutes[2]),"\\\\\n",format(obj$chutes[3]),
		"\n\\end{array} \\right) \\quad \n\\stackunder{\\hat{\\beta}}{\\sim} =\n",
		"\\left(\\begin{array}{l}\n",format(obj$beta[1]),"\\\\\n",
                format(obj$beta[2]),"\\\\\n",format(obj$beta[3]),
                "\n\\end{array} \\right)  \\quad \n\\stackunder{\\hat{\\theta}}{\\sim} = \n",
                "\\left(\\begin{array}{l}\n",format(obj$teta[1]),"\\\\\n",
                format(obj$teta[2]),"\\\\\n",format(obj$teta[3]),"\\\\\n",
		format(obj$teta[4])," \\end{array} \\right) \n\\end{displaymath}\n\n",
		"Estimativas dos valores esperados:\n\n",
                "\\begin{displaymath}\n",
                " \\stackunder{\\hat{\\mu}'}{\\sim} = (",format(round(obj$esperados[1])),",",
                format(round(obj$esperados[2])),",",format(round(obj$esperados[3])),",",
                format(round(obj$esperados[4])),")\n\\end{displaymath}\n\n",
		"\n\nEstatísticas de teste:\n\n",
		"\\begin{center}\n\\begin{tabular}{rrrr}\n\\hline\n & obs & gl & P-valor \\\\ \n \\hline \n ",
		"$Q_V$ & ",format(obj$estatisticas[1,1])," & ",format(obj$estatisticas[1,2])," & ",format(obj$estatisticas[1,3])," \\\\ \n",
                "$Q_P$ & ",format(obj$estatisticas[2,1])," & ",format(obj$estatisticas[2,2])," & ",format(obj$estatisticas[2,3])," \\\\ \n",
                "$Q_N$ & ",format(obj$estatisticas[3,1])," & ",format(obj$estatisticas[3,2])," & ",format(obj$estatisticas[3,3])," \\\\ \n",
		"\\hline\n\\end{tabular}\n\\end{center}\n",
		"\n\\end{minipage}\\vspace{3ex}\n\n",sep="")
  cat(saida,file=file,append=append)		
}

##
## Esse if final serve para definir o método default para a função
## latex. Ele só é executado se ela já não tiver sido definida.
## Ela em geral é definida quando o pacote Hmisc (que define métodos
## latex para outros objetos) está carregado.
##

if (!(exists("latex"))) {
	latex <-
function (object, title = first.word(deparse(substitute(object))), 
    ...) 
 {
    if (!length(oldClass(object))) 
        oldClass(object) <- data.class(object)
    UseMethod("latex")
 }
}