#!/usr/bin/env python
import os, sys
from Spider.Spiderutils import *

"""
Create a doc file of selected particle numbers, and their x,y coordinates.
Given files with particles numbers and coordinates (sndc) and files with
manually selected good particle numbers (good), write out files with
particles numbers and coordinates, for only the good particles.

For particles picked by lfc_pick, results are in coords/sndc*** files, where
x,y coords refer to the center of the image:

 ;spi/hcc   06-OCT-2004 AT 14:03:08   coords/sndc0003.hcc
 ; /         X           Y       PARTICLE NO.   PEAK HT                         
    1 4   1703.0       971.00       1938.0      0.49883     

Manually selected file numbers (using Web) are in good/good*** files:

; hcc/hcc  good/good003.hcc   Tue Sep 28 13:47:33 2004
 ; /           PARTICLE NO.   CLASS                     
 0001 2         1938            1

After renumber.bat:
 ;spi/hcc   28-SEP-2004 AT 15:22:53   good/ngood0003.hcc
 ; /    PARTICLE #                                                              
    1 1   1938.0     
 
"""

def makegoodcoords(sndc, good, outfile):
    " create a doc file of x,y coordinates for good particles"

    X = list2int(readdoc(sndc, column=1))
    Y = list2int(readdoc(sndc, column=2))
    P = list2int(readdoc(sndc, column=3))
    G = list2int(readdoc(good))
    # create a dictionary of coords indexed by particle #
    D = {}
    n = len(P)
    for i in range(n):
        D[P[i]] = (X[i],Y[i])

    datalines=[]
    
    ncol = numberOfColumns(good)
    if ncol == 1:      # use all the particle numbers (good/ngood***)
        for g in G:
            if g in D:
                datalines.append( [D[g][0], D[g][1], g] )

    elif ncol > 1:     # only use particles if class=1   (good/good***)
        C = list2int(readdoc(good, column=2))
        n = len(G)
        for i in range(n):
            if C[i] == 1:
                p = G[i]
                x,y = D[p]
                datalines.append([x,y,p])

    hdrs = ["X","Y","Particle"]
    writedoc(outfile, lines=datalines, headers=hdrs)

def goodcoords(filenumber, dataext):
    fn = int(filenumber)
    if dataext[0] != '.': dataext = '.' + dataext

    sndcf = './coords/sndc****' + dataext   # columns: X, Y, Particle
    goodf = './good/good***' + dataext   # columns: Particle (, class?)
    outf = './cpick/coords****' + dataext

    sndc = makeSpiderFilename(sndcf, fn)
    good = makeSpiderFilename(goodf, fn)
    outfile = makeSpiderFilename(outf, fn)

    print "writing %s" % outfile
    makegoodcoords(sndc, good, outfile)

# ------------------------------------------------------------------------
if __name__ == '__main__':

    nargs = sys.argv[1:]
    if nargs < 1:
        print "Usage: goodcoords.py filenums_docfile"
        sys.exit()
    filenumfile = sys.argv[1]
    
    base,ext = os.path.splitext(filenumfile)  # get the data extension
    filenumbers = readdoc(filenumfile)
    
    for fn in filenumbers:
        goodcoords(fn, ext)