<commit_before>package pdb

import (
	"fmt"

	"github.com/BurntSushi/bcbgo/seq"
)

type Entry struct {
	Path   string
	IdCode string
	Chains []*Chain
}

type Chain struct {
	Entry    *Entry
	Ident    byte
	SeqType  SequenceType
	Sequence []seq.Residue
	Models   []*Model
	Missing  []*Residue
}

type Model struct {
	Entry    *Entry
	Chain    *Chain
	Num      int
	Residues []*Residue
}

type Residue struct {
	Name          seq.Residue
	SequenceNum   int
	InsertionCode byte
	Atoms         []Atom
}

type Atom struct {
	Name string
	Het  bool
	Coords
}

type Coords struct {
	X, Y, Z float64
}

// Chain returns a chain with the given identifier.
// If such a chain does not exist, nil is returned.
func (entry *Entry) Chain(ident byte) *Chain {
	for i := range entry.Chains {
		if entry.Chains[i].Ident == ident {
			return entry.Chains[i]
		}
	}
	return nil
}

// OneChain returns a single chain in the PDB file. If there is more than one
// chain, OneChain will panic. This is convenient when you expect a PDB file to
// have only a single chain, but don't know the name.
func (entry *Entry) OneChain() *Chain {
	if len(entry.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			entry.Path, len(entry.Chains)))
	}
	return entry.Chains[0]
}

// IsProtein returns true if the chain consists of amino acids.
func (c Chain) IsProtein() bool {
	return c.SeqType == SeqProtein && len(c.Models) > 0
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned. If there is more than one model, the first model
// is used.
func (c Chain) SequenceCaAtomSlice(start, end int) []Coords {
	return c.Models[0].SequenceCaAtomSlice(start, end)
}

// SequenceCaAtoms returns a slice of all Ca atoms for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model).
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceCaAtoms() ([]*Coords, error) {
	return c.Models[0].SequenceCaAtoms()
}

// SequenceAtoms returns a slice of all residues for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model). Namely, the mapping is sparse, since not all SEQRES residues have
// an ATOM record.
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceAtoms() ([]*Residue, error) {
	return c.Models[0].SequenceAtoms()
}

// CaAtoms returns all alpha-carbon atoms in the chain. If there is more than
// one model, only the first model is used.
func (c Chain) CaAtoms() []Coords {
	return c.Models[0].CaAtoms()
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned.
func (m Model) SequenceCaAtomSlice(start, end int) []Coords {
	residues, err := m.SequenceCaAtoms()
	if err != nil {
		fmt.Printf("%s\n", err)
		return nil
	}

	atoms := make([]Coords, end-start)
	for i, cai := 0, start; cai < end; i, cai = i+1, cai+1 {
		if residues[start] == nil {
			return nil
		}
		atoms[i] = *residues[cai]
	}
	return atoms
}

// SequenceCaAtoms returns a slice of all Ca atoms for the model in
// correspondence with the sequence in SEQRES.
// Note that a slice of pointers is returned, since not all residues
// necessarily correspond to a alpha-carbon ATOM.
//
// This method can proceed in one of two ways. First, if "REMARK 465" is
// present in the PDB file, it will be used to determine the positions of the
// holes in the sequence (i.e., residues in SEQRES without an ATOM record).
// This method is generally reliable, since REMARK 465 lists all residues
// in SEQRES that don't have an ATOM record. This will fail if there are any
// unreported missing residues.
//
// If "REMARK 465" is absent, then we have to rely on the order of ATOM records
// to correspond to a residue index in the SEQRES sequence. This will fail
// with an error if there are any unreported missing residues.
//<commit_msg>Don't check the same index every time.<commit_after>package pdb

import (
	"fmt"

	"github.com/BurntSushi/bcbgo/seq"
)

type Entry struct {
	Path   string
	IdCode string
	Chains []*Chain
}

type Chain struct {
	Entry    *Entry
	Ident    byte
	SeqType  SequenceType
	Sequence []seq.Residue
	Models   []*Model
	Missing  []*Residue
}

type Model struct {
	Entry    *Entry
	Chain    *Chain
	Num      int
	Residues []*Residue
}

type Residue struct {
	Name          seq.Residue
	SequenceNum   int
	InsertionCode byte
	Atoms         []Atom
}

type Atom struct {
	Name string
	Het  bool
	Coords
}

type Coords struct {
	X, Y, Z float64
}

// Chain returns a chain with the given identifier.
// If such a chain does not exist, nil is returned.
func (entry *Entry) Chain(ident byte) *Chain {
	for i := range entry.Chains {
		if entry.Chains[i].Ident == ident {
			return entry.Chains[i]
		}
	}
	return nil
}

// OneChain returns a single chain in the PDB file. If there is more than one
// chain, OneChain will panic. This is convenient when you expect a PDB file to
// have only a single chain, but don't know the name.
func (entry *Entry) OneChain() *Chain {
	if len(entry.Chains) != 1 {
		panic(fmt.Sprintf("OneChain can only be called on PDB entries with "+
			"ONE chain. But the '%s' PDB entry has %d chains.",
			entry.Path, len(entry.Chains)))
	}
	return entry.Chains[0]
}

// IsProtein returns true if the chain consists of amino acids.
func (c Chain) IsProtein() bool {
	return c.SeqType == SeqProtein && len(c.Models) > 0
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned. If there is more than one model, the first model
// is used.
func (c Chain) SequenceCaAtomSlice(start, end int) []Coords {
	return c.Models[0].SequenceCaAtomSlice(start, end)
}

// SequenceCaAtoms returns a slice of all Ca atoms for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model).
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceCaAtoms() ([]*Coords, error) {
	return c.Models[0].SequenceCaAtoms()
}

// SequenceAtoms returns a slice of all residues for the chain in
// correspondence with the sequence in SEQRES (automatically using the first
// model). Namely, the mapping is sparse, since not all SEQRES residues have
// an ATOM record.
//
// See Model.SequenceCaAtoms for the deets.
func (c Chain) SequenceAtoms() ([]*Residue, error) {
	return c.Models[0].SequenceAtoms()
}

// CaAtoms returns all alpha-carbon atoms in the chain. If there is more than
// one model, only the first model is used.
func (c Chain) CaAtoms() []Coords {
	return c.Models[0].CaAtoms()
}

// SequenceCaAtomSlice attempts to extract a contiguous slice of alpha-carbon
// ATOM records based on *residue* index. Namely, if a contiguous slice cannot
// be found, nil is returned.
func (m Model) SequenceCaAtomSlice(start, end int) []Coords {
	residues, err := m.SequenceCaAtoms()
	if err != nil {
		fmt.Printf("%s\n", err)
		return nil
	}

	atoms := make([]Coords, end-start)
	for i, cai := 0, start; cai < end; i, cai = i+1, cai+1 {
		if residues[cai] == nil {
			return nil
		}
		atoms[i] = *residues[cai]
	}
	return atoms
}

// SequenceCaAtoms returns a slice of all Ca atoms for the model in
// correspondence with the sequence in SEQRES.
// Note that a slice of pointers is returned, since not all residues
// necessarily correspond to a alpha-carbon ATOM.
//
// This method can proceed in one of two ways. First, if "REMARK 465" is
// present in the PDB file, it will be used to determine the positions of the
// holes in the sequence (i.e., residues in SEQRES without an ATOM record).
// This method is generally reliable, since REMARK 465 lists all residues
// in SEQRES that don't have an ATOM record. This will fail if there are any
// unreported missing residues.
//
// If "REMARK 465" is absent, then we have to rely on the order of ATOM records
// to correspond to a residue index in the SEQRES sequence. This will fail
// with an error if there are any unreported missing residues.
//
{'fraction_non_alphanumeric': 0.09777242624924744, 'fraction_numerical': 0.005057194461167971, 'mean_word_length': 3.399364406779661, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3342170', 'n_tokens_mistral': 2798, 'n_tokens_neox': 2604, 'n_words': 1136}