MTAutocolors/cmd/mtautocolors/main.go

// Copyright 2015 by Sascha L. Teichmann
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package main

import (
	"bufio"
	"flag"
	"fmt"
	"image/color"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"
	"sync"
)

var nodesSplit = regexp.MustCompile(" +")

func usage() {
	fmt.Fprintf(os.Stderr,
		"Usage: %s [<options>] <nodes.txt> [<data source directory> ...]\n",
		os.Args[0])
	fmt.Fprintln(os.Stderr, "Options:")
	flag.PrintDefaults()
}

func main() {
	var predef string
	var transparent string
	var workers int

	flag.Usage = usage
	flag.StringVar(&predef, "predefined", "", "predefined colors")
	flag.StringVar(&predef, "p", "", "predefined colors (shorthand)")
	flag.StringVar(&transparent, "transparent", "glasslike,liquid", "transparent nodes")
	flag.StringVar(&transparent, "t", "glasslike,liquid", "transparent nodes (shorthand)")
	flag.IntVar(&workers, "workers", 0, "number of image processing workers")
	flag.IntVar(&workers, "w", 0, "number of image processing workers (shorthand)")
	flag.Parse()

	nargs := flag.NArg()
	if nargs < 1 {
		usage()
		os.Exit(1)
	}

	var err error
	var predefs PredefCols

	if predef != "" {
		if predefs, err = LoadPredefCols(predef); err != nil {
			log.Fatalf("Cannot load predefined colors: %s\n", err)
		}
	}

	var roots []string
	if nargs > 1 {
		roots = flag.Args()[1:]
	} else {
		roots = []string{"."}
	}

	files, err := buildFileIndex(roots)
	if err != nil {
		log.Fatalf("error while building file index: %s\n", err)
	}

	drawTypes := strings.Split(transparent, ",")

	if err = process(flag.Arg(0), drawTypes, files, predefs, workers); err != nil {
		log.Fatalf("error while generating colors: %s\n", err)
	}
}

func contains(haystack []string, needle string) bool {
	for _, straw := range haystack {
		if strings.Contains(needle, straw) {
			return true
		}
	}
	return false
}

type outLine struct {
	name  string
	col   color.Color
	alpha bool
}

func asByte(x uint32) byte {
	return byte(x >> 8)
}

func (ol *outLine) print(out io.Writer) {
	r, g, b, a := ol.col.RGBA()
	ba := asByte(a)
	if ol.alpha && ba < 255 {
		fmt.Fprintf(out, "%s %d %d %d %d\n",
			ol.name, asByte(r), asByte(g), asByte(b), ba)
	} else {
		fmt.Fprintf(out, "%s %d %d %d\n",
			ol.name, asByte(r), asByte(g), asByte(b))
	}
}

func process(
	nodesFile string,
	drawTypes []string,
	files map[string]string,
	predefs PredefCols,
	workers int) error {

	file, err := os.Open(nodesFile)
	if err != nil {
		return err
	}
	defer file.Close()

	lineCh := make(chan []string)

	if workers < 1 {
		workers = runtime.NumCPU()
	}

	var outLineMu sync.Mutex
	var outLines []outLine

	var wg sync.WaitGroup

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
		lines:
			for line := range lineCh {
				name, drawType, texture := line[0], line[1], line[2]

				var c color.Color
				switch col := predefs.findCol(name); {
				case col != nil && col.Complete():
					c = col
				case col == nil || !col.Complete():
					tfile := files[texture]
					if tfile == "" {
						log.Printf("WARN: node '%s' missing texture '%s'\n", name, texture)
						continue lines
					}
					avg, err := averageColor(tfile)
					if err != nil {
						log.Printf("WARN: node '%s' defect image: %s\n", name, err)
						continue lines
					}
					if col != nil {
						c = col.Apply(avg)
					} else {
						c = avg
					}
				}
				alpha := contains(drawTypes, drawType)
				outLineMu.Lock()
				outLines = append(outLines, outLine{
					name:  name,
					col:   c,
					alpha: alpha,
				})
				outLineMu.Unlock()
			}
		}()
	}

	scanner := bufio.NewScanner(file)
	for lineNo := 1; scanner.Scan(); lineNo++ {
		parts := nodesSplit.Split(scanner.Text(), 3)
		if len(parts) < 3 {
			log.Printf("WARN: line %d too short.\n", lineNo)
		} else {
			lineCh <- parts
		}
	}
	close(lineCh)
	wg.Wait()
	if err := scanner.Err(); err != nil {
		return err
	}

	// To make it more deterministic.
	sort.Slice(outLines, func(i, j int) bool {
		return outLines[i].name < outLines[j].name
	})

	out := bufio.NewWriter(os.Stdout)
	for i := range outLines {
		outLines[i].print(out)
	}
	return out.Flush()
}

func buildFileIndex(roots []string) (map[string]string, error) {
	index := make(map[string]string)

	acceptedExts := map[string]bool{
		".png":  true,
		".jpg":  true,
		".jpeg": true}

	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.Mode().IsRegular() {
			return nil
		}
		name := info.Name()
		if !acceptedExts[strings.ToLower(filepath.Ext(name))] {
			return nil
		}
		if _, found := index[name]; !found {
			index[name] = path
		} else {
			log.Printf("WARN: more than one file for '%s'\n", name)
		}
		return nil
	}

	for _, root := range roots {
		if err := filepath.Walk(root, walkFn); err != nil {
			return nil, err
		}
	}
	return index, nil
}