<commit_before>package main

import (
	"fmt"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/los/tree"
)

type IDType int
const (
	Rockstar IDType = iota
	M500c
	M200c
	M200m

    finderCells = 150
    overlapMult = 3
)

var (
	ms, rs, xs, ys, zs []float64
	rids []int
)

func main() {
	// Parse command line
	var (
		idTypeStr string
		idStart, idEnd, snap int
		allowSubhalos bool
	)

	flag.StringVar(&idTypeStr, "IDType", "Rockstar",
		"[Rockstar | M500c | M200c | M200m]")
	flag.IntVar(&snap, "Snap", -1, "The snapshot number.")
	flag.IntVar(&idStart, "IDStart", -1, "ID start range (inclusive).")
	flag.IntVar(&idEnd, "IDEnd", -1, "ID stop range (exclusive)")
	flag.BoolVar(&allowSubhalos, "AllowSubhalos", false,
		"Allow subhalo ids to be passed through.")
	flag.Parse()

	if idStart > idEnd {
		log.Fatalf("IDStart, %d, is larger than IDEnd, %d.")
	} else if idStart != idEnd && idStart <= 0 {
		log.Fatalf("Non-positive IDStart %d.")
	}

	idType, err := parseIDType(idTypeStr)
	if err != nil { err.Error() }
	if idType != Rockstar && snap == -1 {
		log.Fatalf("Must set the Snap flag if using a non-default IDType.")
	}

	snapNum, err := snapNum()
	if err != nil {
		log.Fatalf(
			"Error encountered when finding rockstar directory: %s",err.Error(),
		)
	} else if (snap < 1 && idType != Rockstar) || snap > snapNum {
		log.Fatalf("Snap %d is out of bounds for %d snaps.", snap, snapNum)
	}

	// Get IDs and snapshots

	rawIds := getIDs(idStart, idEnd, flag.Args())
	if len(rawIds) == 0 { return }

	fmt.Println(rawIds)

	var ids, snaps []int
	switch idType {
	case Rockstar:
		if snap != -1 {
			snaps = make([]int, len(rawIds))
			for i := range snaps { snaps[i] = snap }
		} else {
			snaps, err = findSnaps(rawIds)
		}
		ids = rawIds
		if err != nil { log.Fatalf(err.Error()) }
	case M200m:
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = snap }
		ids, err = convertSortedIDs(rawIds, snap)
		if err != nil { err.Error() }
	default:
		log.Fatal("Unsupported IDType for now. Sorry :3")
	}

	// Tag subhalos, if neccessary.
	var isSub []bool
	if allowSubhalos {
		isSub = make([]bool, len(ids))
	} else {
		isSub, err = findSubs(rawIds, snaps)
		if err != nil { log.Fatal(err.Error()) }
	}

	// Output
	printIds(ids, snaps, isSub)
}

func parseIDType(str string) (IDType, error) {
	switch strings.ToLower(str) {
	case "rockstar": return Rockstar, nil
	case "m500c": return M500c, nil
	case "m200c": return M200c, nil
	case "m200m": return M200m, nil
	}
	return -1, fmt.Errorf("IDType '%s' unrecognized", str)
}

func getIDs(idStart, idEnd int, args []string) []int {
	ids := make([]int, 0, idEnd - idStart + len(args))
	for i := idStart; i < idEnd; i++ { ids = append(ids, i) }
	for _, str := range args {
		i, err := strconv.Atoi(str)
		if err != nil {
			log.Fatalf("Could not parse arg %d: '%s' is not an int.", i+1, str)
		}
		ids = append(ids, i)
	}
	return ids
}

func snapNum() (int, error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	return len(infos), err
}

func getSnapHaloList(i int) (name string, err error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	if err != nil { return "", err }
	return path.Join(rockstarDir, infos[i - 1].Name()), nil
}

func findSnaps(ids []int) ([]int, error) {
	treeDir := os.Getenv("GTET_TREE_DIR")
	infos, err := ioutil.ReadDir(treeDir)
	if err != nil { return nil, err }

	names := []string{}
	for _, info := range infos {
		name := info.Name()
		n := len(name)
		if n > 4 && name[:5] == "tree_" && name[n-4:] == ".gtet" {
			names = append(names, name)
		}
	}

	return tree.HaloSnaps(names, ids)
}


// This funciton has side effects :(
func convertSortedIDs(
	rawIDs []int, snap int,
) ([]int, error) {
	list, err := getSnapHaloList(snap)
	if err != nil { return nil, err }

	gtetFmt := os.Getenv("GTET_FMT")
	gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(gtetName, hd)
	if err != nil { return nil, err }
	cosmo := &hd.Cosmo

	rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(list, halo.R200m, cosmo)
	if err != nil { return nil, err }
	
	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil<commit_msg>Fixed bug in tree_*.dat searching code.<commit_after>package main

import (
	"fmt"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/los/tree"
)

type IDType int
const (
	Rockstar IDType = iota
	M500c
	M200c
	M200m

    finderCells = 150
    overlapMult = 3
)

var (
	ms, rs, xs, ys, zs []float64
	rids []int
)

func main() {
	// Parse command line
	var (
		idTypeStr string
		idStart, idEnd, snap int
		allowSubhalos bool
	)

	flag.StringVar(&idTypeStr, "IDType", "Rockstar",
		"[Rockstar | M500c | M200c | M200m]")
	flag.IntVar(&snap, "Snap", -1, "The snapshot number.")
	flag.IntVar(&idStart, "IDStart", -1, "ID start range (inclusive).")
	flag.IntVar(&idEnd, "IDEnd", -1, "ID stop range (exclusive)")
	flag.BoolVar(&allowSubhalos, "AllowSubhalos", false,
		"Allow subhalo ids to be passed through.")
	flag.Parse()

	if idStart > idEnd {
		log.Fatalf("IDStart, %d, is larger than IDEnd, %d.")
	} else if idStart != idEnd && idStart <= 0 {
		log.Fatalf("Non-positive IDStart %d.")
	}

	idType, err := parseIDType(idTypeStr)
	if err != nil { err.Error() }
	if idType != Rockstar && snap == -1 {
		log.Fatalf("Must set the Snap flag if using a non-default IDType.")
	}

	snapNum, err := snapNum()
	if err != nil {
		log.Fatalf(
			"Error encountered when finding rockstar directory: %s",err.Error(),
		)
	} else if (snap < 1 && idType != Rockstar) || snap > snapNum {
		log.Fatalf("Snap %d is out of bounds for %d snaps.", snap, snapNum)
	}

	// Get IDs and snapshots

	rawIds := getIDs(idStart, idEnd, flag.Args())
	if len(rawIds) == 0 { return }

	var ids, snaps []int
	switch idType {
	case Rockstar:
		if snap != -1 {
			snaps = make([]int, len(rawIds))
			for i := range snaps { snaps[i] = snap }
		} else {
			snaps, err = findSnaps(rawIds)
		}
		ids = rawIds
		if err != nil { log.Fatalf(err.Error()) }
	case M200m:
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = snap }
		ids, err = convertSortedIDs(rawIds, snap)
		if err != nil { err.Error() }
	default:
		log.Fatal("Unsupported IDType for now. Sorry :3")
	}

	// Tag subhalos, if neccessary.
	var isSub []bool
	if allowSubhalos {
		isSub = make([]bool, len(ids))
	} else {
		isSub, err = findSubs(rawIds, snaps)
		if err != nil { log.Fatal(err.Error()) }
	}

	// Output
	printIds(ids, snaps, isSub)
}

func parseIDType(str string) (IDType, error) {
	switch strings.ToLower(str) {
	case "rockstar": return Rockstar, nil
	case "m500c": return M500c, nil
	case "m200c": return M200c, nil
	case "m200m": return M200m, nil
	}
	return -1, fmt.Errorf("IDType '%s' unrecognized", str)
}

func getIDs(idStart, idEnd int, args []string) []int {
	ids := make([]int, 0, idEnd - idStart + len(args))
	for i := idStart; i < idEnd; i++ { ids = append(ids, i) }
	for _, str := range args {
		i, err := strconv.Atoi(str)
		if err != nil {
			log.Fatalf("Could not parse arg %d: '%s' is not an int.", i+1, str)
		}
		ids = append(ids, i)
	}
	return ids
}

func snapNum() (int, error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	return len(infos), err
}

func getSnapHaloList(i int) (name string, err error) {
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	infos, err := ioutil.ReadDir(rockstarDir)
	if err != nil { return "", err }
	return path.Join(rockstarDir, infos[i - 1].Name()), nil
}

func findSnaps(ids []int) ([]int, error) {
	treeDir := os.Getenv("GTET_TREE_DIR")
	infos, err := ioutil.ReadDir(treeDir)
	if err != nil { return nil, err }

	names := []string{}
	for _, info := range infos {
		name := info.Name()
		n := len(name)
		if n > 4 && name[:5] == "tree_" && name[n-4:] == ".dat" {
			names = append(names, path.Join(treeDir, name))
		}
	}

	return tree.HaloSnaps(names, ids)
}


// This funciton has side effects :(
func convertSortedIDs(
	rawIDs []int, snap int,
) ([]int, error) {
	list, err := getSnapHaloList(snap)
	if err != nil { return nil, err }

	gtetFmt := os.Getenv("GTET_FMT")
	gtetName := fmt.Sprintf(gtetFmt, snap, 0, 0, 0)
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(gtetName, hd)
	if err != nil { return nil, err }
	cosmo := &hd.Cosmo

	rids, xs, ys, zs, ms, rs, err = halo.ReadRockstar(list, halo.R200m, cosmo)
	if err != nil { return nil, err }
	
	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil
{'fraction_non_alphanumeric': 0.16570978273105932, 'fraction_numerical': 0.015084993808398063, 'mean_word_length': 3.6174636174636174, 'pattern_counts': {'":': 8, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '311953', 'n_tokens_mistral': 3788, 'n_tokens_neox': 3530, 'n_words': 970}