# Copyright (c) 1996, 1997, 1998 The Regents of the University of California.
# All rights reserved.  See legal notices for full text and disclaimer.

from history_medium import *
from history import *
from PW import PW

class PDBfile (HistoryMedium):
    "A history medium using a PDB file."
    verbose = 0

    def __init__ (self, file_name, family_size = 0):
        """Create PDBfile medium with given base filename and family
        size. A family size of 0 means do not family the files."""

        HistoryMedium.initialize (self)
        self.set_name (file_name)
        self.family_size = family_size  # number of records to allow per file
        self.family_count = 0 # number of files in the family so far

    def __repr__ (self):
        return "PDBfile History Medium, file = " + self.filename ()

    __str__ = __repr__

    def set_name (self, file_name):
        "Set file name to file_name."
        print "set_name : " + `file_name`
        if self.is_open:
            raise RunTimeError, \
                  "Cannot change filename while record is open."
        
        if file_name [-4 : ] == ".pdb":
            self._filename_base = file_name [0 : -4]
        else:
            self._filename_base = file_name

        self.file_record_number = 0

    def filename (self):
        if self.family_size:
            count = self.family_count
            if count < 10:
                return self._filename_base + ".00" + `count` + ".pdb"
            if count < 100:
                return self._filename_base + ".0" + `count` + ".pdb"
            return self._filename_base + "." + `count` + ".pdb"
        else:
            return self._filename_base + ".pdb"

    def begin_record (self, cycle, time):
        if self.family_size and self.file_record_number >= self.family_size:
            self.family_count = self.family_count + 1
            self.file_record_number = 0
        HistoryMedium.begin_record (self, cycle, time)
        f = self.filename ()
        if PDBfile.verbose:
            print \
                  "File", f, \
                  ", record number = ", self.record_number,\
                  ", cycle =", cycle, \
                  ", time = ", time

        if self.file_record_number == 0:
            self.writer = PW (f)
            if PDBfile.verbose:
                print "PDB history created " + f
        else:
            self.writer = PW ()
            self.writer.open (f, 'a')

        self.file_record_number = self.file_record_number + 1

        if PDBfile.verbose > 1:
            self.writer.set_verbosity (PDBfile.verbose)
        else:
            self.writer.set_verbosity (0)

        self.writer.write ("cycle", cycle, self.file_record_number)
        self.writer.write ("time", time, self.file_record_number)

    def end_record (self):
        self.writer.close ()
        self.writer = None
        HistoryMedium.end_record (self)
        if PDBfile.verbose:
            print "End record", self.record_number

    def _write (self, ename, value):
        self.writer.write (ename, value, self.file_record_number)

def PDBfile_tag (name, file_name = None, family_size = 0, col = collector):
    "Create a tag based on a PDB file."
    if file_name is None:
        f = name 
    else:
        f = file_name
    return Tag(name, PDBfile (f, family_size), col)

if __name__ == "__main__":
    h = PDBfile_tag ("pdbtag1")
    k = PDBfile_tag ("pdbtag2", "tag2_history.pdb", 5)
    PDBfile.verbose = 2
    h.item ("x")
    h.item ("y", "x*x")
    h.item_static ("s", "'hello world'")
    h.frequency (10, 50, 5)
    k.item ("x")

    for i in range (100):
        x = .2 * i
        time = i/10.
        cycle = i
        collector.sample (cycle, time)
    collector.sample_final (cycle, time)
    raise SystemExit